from typing import Literal

import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa
from keras import backend as K
from scipy.spatial.distance import directed_hausdorff
from tensorflow import keras

CLASSES = ["background", "media", "lumen"]
# Classification metrics

categorical_crossentropy = keras.losses.CategoricalCrossentropy()

# calculates dice considering an input with a single class


def dice_single(true, pred, smooth=1.0):
    true = K.batch_flatten(true)
    pred = K.batch_flatten(pred)
    pred = K.round(pred)

    intersection = K.sum(true * pred, axis=-1)
    true = K.sum(true, axis=-1)
    pred = K.sum(pred, axis=-1)

    return ((2 * intersection) + smooth) / (true + pred + smooth)


def dice_for_class(index, smooth=1.0):
    """Evaluates dice coefficient for a given segmentation class"""
    def dice_inner(true, pred):

        # get only the desired class
        true = true[:, :, :, index]
        pred = pred[:, :, :, index]

        # return dice per class
        return dice_single(true, pred, smooth)

    dice_inner.__name__ = f"dice_{CLASSES[index]}"
    return dice_inner


def dice_coef_loss(y_true, y_pred, smooth=1.0):
    """Dice coefficient-based loss for image segmentation

    Original implementation by Albert Sallés based on [1]_.

    .. [1] https://arxiv.org/pdf/1606.04797v1.pdf
    """
    # Ignore label 0 which belongs to background
    y_true = K.flatten(y_true)
    y_pred = K.flatten(y_pred)
    intersection = K.sum(y_true * y_pred, axis=-1)
    denom = K.sum(y_true + y_pred, axis=-1)
    dice = K.mean(((2.0 * intersection + smooth) / (denom + smooth)))
    return 1 - dice


def jaccard_coef_loss(y_true, y_pred, smooth=1.0):
    """Jaccard distance for semantic segmentation.

    Also known as the intersection-over-union loss. Originally implemented by
    Albert Sallès using [1]_.

    Parameters
    -----------
    y_true : Tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    y_pred : Tensor
        The target distribution, format the same with `output`.
    smooth : float
        This small value will be added to the numerator and denominator.
            - If both output and target are empty, it makes sure dice is 1.
            - If either output or target are empty (all pixels are background),
            dice = ```smooth/(small_value + smooth)``,
            then if smooth is very small, dice close to 0 (even the image values lower than the threshold),
            so in this case, higher smooth can have a higher dice.

    .. [1] https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/jaccard.py
    """
    intersection = K.sum(K.abs(y_true * y_pred), axis=-1)
    sum_ = K.sum(K.abs(y_true) + K.abs(y_pred), axis=-1)
    jac = (intersection + smooth) / (sum_ - intersection + smooth)
    return (1 - jac) * smooth


def jaccard_index(y_true, y_pred):
    ax = (-2, -1)
    p = np.round(y_true)

    q = np.round(y_pred)

    intersection = np.sum(np.abs(p * q), ax)
    union = np.sum(p, ax) + np.sum(q, ax) - intersection
    return intersection / union


def round_diff(x, k=1e7):
    return tf.sigmoid(k * (x - 0.5))


def CombinedLoss(smooth=1.0, hausdorff_weight=1.0 / 3, weights=None):
    """Combine both `JaccardLoss` and `HausdorffDTLoss` into a loss 

    Parameters
    ----------
    smooth : float 
        small value used as hyperparameter for the jaccard loss part 
    hausdorff_weight : float 
        value that controls how much the hausdorff distance (approximation)
        influences the result
    weights : numpy.ndarray -- with shape (1,) or (3,)
        used as the weights for each layer when combining them. The default
        value is `numpy.ndarray([1.0])`, i.e. all layers have the same weight.
    """
    if weights is None:
        weights = np.array([1.0])

    def inner(true, pred):
        """jaccard"""
        inter = K.sum(true * pred, (-3, -2))
        union = K.sum(true, (-3, -2)) + K.sum(pred, (-3, -2)) - inter
        jacc = (inter + smooth) / (union + smooth)
        jacc = K.sum((1 - jacc) * smooth, -1)

        """Hausdorff"""
        p, q = true, pred
        dp = tfa.image.euclidean_dist_transform(
            1 - tf.cast(p, tf.uint8), dtype=p.dtype)
        dp = dp + p  # ensures that the los is 0 iff p = q

        diff = K.pow(p - q, 2)
        dist = K.pow(dp, 2.0)  # + dq
        hd = K.mean(diff * dist, axis=(-3, -2)) * weights
        hd = K.sum(hd, axis=-1)

        """combined"""
        return jacc + K.sqrt(hd) * hausdorff_weight

    return inner


def HausdorffDTLoss(alpha=2.0,
                    crit: Literal["max", "sum"] = "sum",
                    weights=None):
    """Approximage the Hausdorff distance using distance transforms
    as seen in [1]_

    Parameters
    ----------
    alpha : float  
        hyperparameter
    crit  : "max", "sum"
        how to combine layers

    .. [1] D. Karimi, S. E. Salcudean "Reducing the Hausdorff Distance in
           Medical Image Segmentation with Convolutional Neural Networks"
           https://arxiv.org/abs/1904.10030
    """

    if crit == "max":
        reducer = tf.math.reduce_max
    elif crit == "sum":
        reducer = tf.math.reduce_sum
    else:
        raise NotImplementedError(crit)

    if weights is not None:
        weights = weights / np.sum(weights)
    else:
        # Weights obtained empirically by measuring the relative size of the sets
        weights = np.array([1.0])

    def inner(truth, predicted):
        p, q = truth, predicted
        dp = tfa.image.euclidean_dist_transform(1 - tf.cast(p, tf.uint8),
                                                dtype=p.dtype)
        dp = dp + p  # ensures that the los is 0 iff p = q
        diff = K.pow(p - q, 2)
        dist = K.pow(dp, alpha)
        result = K.mean(diff * dist, axis=(-3, -2)) * weights
        return reducer(result, axis=-1)

    return inner


def hausdorff_distance(y_true, y_pred):
    r"""Calculate the Hausdorff distance between two sets. 

    Parameters
    ----------
    y_true : ndarray -- shape must be (N, M)
    y_pred : ndarray -- shape must be (N, S)

    Notes
    -----
    .. math:: \text{HD}(X, Y) = \max(\text{hd}(X, Y), \text{hd}(Y, X))
    .. math:: \text{hd}(X, Y) = \min_{x \in X} max{ y \in Y } |x-y|
    """
    P = np.stack(y_true.nonzero(), 1)
    Q = np.stack(y_pred.nonzero(), 1)
    d1 = directed_hausdorff(P, Q)[0]
    d2 = directed_hausdorff(Q, P)[0]
    return max(d1, d2)
