import logging

import GeodisTK
import numpy as np

logger = logging.getLogger(__name__)


def geodesic_distance_2d(I, S, lamda=1.0, iter=2):
    # lamda=0 : euclidean dist
    # lamda=1 : geodesic dist
    return GeodisTK.geodesic2d_raster_scan(I, S, lamda, iter)


def geodesic_distance_3d(I, S, spacing, lamda=1.0, iter=4):
    # lamda=0 : euclidean dist
    # lamda=1 : geodesic dist
    return GeodisTK.geodesic3d_raster_scan(I, S, spacing, lamda, iter)


def stable_softmax_foreground_return(x):
    # expects first dimension to have backgroundidx=0, foregroundidx=1
    # returns only worked out values for foreground, background can be reconstructed
    # as 1-Foreground in the calling function
    z = x - np.max(x, axis=0)
    numerator = np.exp(z)
    denominator = np.sum(numerator, axis=0)
    # only return for foreground
    return numerator[1, ...] / denominator


def mideepegd(D_fg, D_bg, P_fg, tau=1):
    """
    Implementing Equation 3 to 7 from Section 2.4 page 6 of: https://arxiv.org/pdf/2104.12166.pdf
    Additionally adding temperature term (\tau) to exponential in Equation 3 and 4
    For replicating papers implementation use \tau==1
    """
    # do ExponentialGaussianDistance calculation
    D_g = np.array([D_bg, D_fg])
    E_fg = stable_softmax_foreground_return(-D_g / tau)

    # calculate alpha_i for each element to be updated
    # Equation 7
    D_g_min = np.min(D_g, axis=0)
    alpha_i = np.exp(-D_g_min)

    # update prob using alpha_i
    R_fg = (1 - alpha_i) * P_fg + alpha_i * E_fg
    return np.array([1 - R_fg, R_fg])


def make_mideepnd_unary(
    image,
    prob,
    scribbles,
    scribbles_fg_label,
    scribbles_bg_label,
    spacing,
    tau=1,
):
    # inputs are expected to be of format [1, X, Y, [Z]]
    # simple check to see if input shape is expected
    if image.shape[0] != 1 or scribbles.shape[0] != image.shape[0]:
        raise ValueError("unexpected input shape received")

    # only binary probabilities supported at the moment
    if prob.shape[0] != 2:
        raise ValueError("only binary probabilities are supported at the moment, received {}".format(prob.shape[0]))

    # extract spatial dims
    spatial_dims = image.ndim - 1
    run_3d = spatial_dims == 3

    P_fg = prob[1, ...]
    scribbles = np.squeeze(scribbles)
    image = np.squeeze(image)

    # get Geodesic Distances for foreground (D_fg) and background (D_bg) scribbles
    S_fg = (scribbles == scribbles_fg_label).astype(np.uint8)
    S_bg = (scribbles == scribbles_bg_label).astype(np.uint8)

    if run_3d:
        D_fg = geodesic_distance_3d(image, S_fg, spacing, lamda=1.0)
        D_bg = geodesic_distance_3d(image, S_bg, spacing, lamda=1.0)
    else:
        D_fg = geodesic_distance_2d(image, S_fg, lamda=1.0)
        D_bg = geodesic_distance_2d(image, S_bg, lamda=1.0)

    # run Equation 3 to 7
    unary = mideepegd(D_fg, D_bg, P_fg, tau=tau)  # \tau=1 for paper default

    return unary


import logging

import GeodisTK
import numpy as np

logger = logging.getLogger(__name__)


def geodesic_distance_2d(I, S, lamda=1.0, iter=2):
    # lamda=0 : euclidean dist
    # lamda=1 : geodesic dist
    return GeodisTK.geodesic2d_raster_scan(I, S, lamda, iter)


def geodesic_distance_3d(I, S, spacing, lamda=1.0, iter=4):
    # lamda=0 : euclidean dist
    # lamda=1 : geodesic dist
    return GeodisTK.geodesic3d_raster_scan(I, S, spacing, lamda, iter)


def stable_softmax_foreground_return(x):
    # expects first dimension to have backgroundidx=0, foregroundidx=1
    # returns only worked out values for foreground, background can be reconstructed
    # as 1-Foreground in the calling function
    z = x - np.max(x, axis=0)
    numerator = np.exp(z)
    denominator = np.sum(numerator, axis=0)
    # only return for foreground
    return numerator[1, ...] / denominator


def mideepegd(D_fg, D_bg, P_fg, tau=1):
    """
    Implementing Equation 3 to 7 from Section 2.4 page 6 of: https://arxiv.org/pdf/2104.12166.pdf
    Additionally adding temperature term (\tau) to exponential in Equation 3 and 4
    For replicating papers implementation use \tau==1
    """
    # do ExponentialGaussianDistance calculation
    D_g = np.array([D_bg, D_fg])
    E_fg = stable_softmax_foreground_return(-D_g / tau)

    # calculate alpha_i for each element to be updated
    # Equation 7
    D_g_min = np.min(D_g, axis=0)
    alpha_i = np.exp(-D_g_min)

    # update prob using alpha_i
    R_fg = (1 - alpha_i) * P_fg + alpha_i * E_fg
    return np.array([1 - R_fg, R_fg])


def make_mideepnd_unary(
    image,
    prob,
    scribbles,
    scribbles_fg_label,
    scribbles_bg_label,
    spacing,
    tau=1,
):
    # inputs are expected to be of format [1, X, Y, [Z]]
    # simple check to see if input shape is expected
    if image.shape[0] != 1 or scribbles.shape[0] != image.shape[0]:
        raise ValueError("unexpected input shape received")

    # only binary probabilities supported at the moment
    if prob.shape[0] != 2:
        raise ValueError("only binary probabilities are supported at the moment, received {}".format(prob.shape[0]))

    # extract spatial dims
    spatial_dims = image.ndim - 1
    run_3d = spatial_dims == 3

    P_fg = prob[1, ...]
    scribbles = np.squeeze(scribbles)
    image = np.squeeze(image)

    # get Geodesic Distances for foreground (D_fg) and background (D_bg) scribbles
    S_fg = (scribbles == scribbles_fg_label).astype(np.uint8)
    S_bg = (scribbles == scribbles_bg_label).astype(np.uint8)

    if run_3d:
        D_fg = geodesic_distance_3d(image, S_fg, spacing, lamda=1.0)
        D_bg = geodesic_distance_3d(image, S_bg, spacing, lamda=1.0)
    else:
        D_fg = geodesic_distance_2d(image, S_fg, lamda=1.0)
        D_bg = geodesic_distance_2d(image, S_bg, lamda=1.0)

    # run Equation 3 to 7
    unary = mideepegd(D_fg, D_bg, P_fg, tau=tau)  # \tau=1 for paper default

    return unary


# lamda=0 : euclidean dist # lamda=1 : geodesic dist # lamda=0 : euclidean dist # lamda=1 : geodesic dist # expects first dimension to have backgroundidx=0, foregroundidx=1 # returns only worked out values for foreground, background can be reconstructed # as 1-Foreground in the calling function # only return for foreground Implementing Equation 3 to 7 from Section 2.4 page 6 of: https://arxiv.org/pdf/2104.12166.pdf     Additionally adding temperature term (\tau) to exponential in Equation 3 and 4     For replicating papers implementation use \tau==1 # do ExponentialGaussianDistance calculation # calculate alpha_i for each element to be updated # Equation 7 # update prob using alpha_i # inputs are expected to be of format [1, X, Y, [Z]] # simple check to see if input shape is expected # only binary probabilities supported at the moment # extract spatial dims # get Geodesic Distances for foreground (D_fg) and background (D_bg) scribbles # run Equation 3 to 7 # \tau=1 for paper default