"""" Script to compute different loss functions in Keras based on tensorflow.

This script compute dice loss, binary cross entropy loss, focal loss, and their combinations.

It also computes hard and soft dice metric as well as loss.

"""

# Import libraries

from numpy.random import seed

from typing import List, Tuple

from keras import backend as K

import tensorflow as tf

from numpy import ndarray

# seed random number generator

seed(1)

class LossMetric:

    """ compute loss and metrics

    Attributes:

        y_true: the reference value,

        y_predicted: the predicted value to compare with y_true.

    Returns:

        Returns the loss or metric.

    """

    def __init__(self, y_true: List[float] = None, y_predicted: List[float] = None):

        self.y_true = y_true

        self.y_predicted = y_predicted

    @staticmethod

    def dice_metric(y_true: ndarray = None, y_predicted: ndarray = None, soft_dice: bool = False,

            threshold_value: float = 0.5, smooth=1) -> float:

        """compute the dice coefficient between the reference and target

        Threshold dice similarity coefficient

        Args:

            y_true: reference target.

            y_predicted:  predicted target by the model.

            soft_dice: apply soft dice or not.

            threshold_value:  thresholding value for soft-dice application.

            smooth: avoid division by zero values.

        Returns:

            Returns dice similarity coefficient, with threshold predicted values

        """

        y_true = K.flatten(y_true)

        y_predicted = K.flatten(y_predicted)

        # prevent from log(0)

        y_true = K.clip(y_true, 10e-8, 1. - 10e-8)

        y_predicted = K.clip(y_predicted, 10e-8, 1. - 10e-8)

        # soft dice

        if soft_dice:

            y_predicted = K.cast(K.greater(y_predicted, threshold_value), dtype='float32')

        intersection = K.sum(y_true * y_predicted)

        # smooth: avoid by zero division

        dice = (2. * intersection + smooth) / (K.sum(y_true) + K.sum(y_predicted) + smooth)

        return dice

    def dice_loss(self, y_true: ndarray, y_predicted: ndarray) -> float:

        """ Compute the dice loss

        Args:

            y_true: reference target.

            y_predicted:  predicted target by the model.

        Returns:

            Returns dice loss.

        """

        return 1 - self.dice_metric(y_true, y_predicted)

    @staticmethod

    def binary_cross_entropy_loss(y_true: ndarray = None, y_predicted: ndarray = None) -> float:

        """ compute the binary cross entropy loss

        Args:

            y_true: reference target.

            y_predicted:  predicted target by the model.

        Returns:

            Returns binary cross entropy between the target and predicted value.

        """

        # prevent from log(0)

        y_true = K.clip(y_true, 10e-8, 1. - 10e-8)

        y_predicted = K.clip(y_predicted, 10e-8, 1. - 10e-8)

        return K.binary_crossentropy(y_true, y_predicted)

    @staticmethod

    def binary_focal_loss(y_true: ndarray = None, y_predicted: ndarray = None, gamma: int = 2,

            alpha: float = .25) -> float:

        """ computes the focal loss

        Args:

            y_true: reference target.

            y_predicted: predicted target by the model.

            gamma: constant value

            alpha: constant value

        Returns:

            Returns focal loss.

        """

        y_true = K.cast(y_true, dtype='float32')

        # Define epsilon so that the back-propagation will not result in NaN for 0 divisor case

        epsilon = K.epsilon()

        # Add the epsilon to prediction value

        y_predicted = y_predicted + epsilon

        # Clip the prediction value

        y_predicted = K.clip(y_predicted, epsilon, 1.0 - epsilon)

        # Calculate p_t

        p_t = tf.where(K.equal(y_true, 1), y_predicted, 1 - y_predicted)

        # Calculate alpha_t

        alpha_factor = K.ones_like(y_true) * alpha

        alpha_t = tf.where(K.equal(y_true, 1), alpha_factor, 1 - alpha_factor)

        # Calculate cross entropy

        cross_entropy = -K.log(p_t)

        weight = alpha_t * K.pow((1 - p_t), gamma)

        # Calculate focal loss f

        loss = weight * cross_entropy

        # Sum the losses in mini_batch

        loss = K.mean(K.sum(loss, axis=1))

        return loss

    @staticmethod

    def focal_loss(y_true: ndarray = None, y_predicted: ndarray = None, gamma: int = 2, alpha: float = .25) -> float:

        """ computes the focal loss

        Adapted from: https://www.kaggle.com/code/bigironsphere/loss-function-library-keras-pytorch/notebook

        """

        y_true = K.flatten(y_true)

        y_predicted = K.flatten(y_predicted)

        bce = K.binary_crossentropy(y_true, y_predicted)

        bce_exp = K.exp(-bce)

        focal_loss = K.mean(alpha * K.pow((1 - bce_exp), gamma) * bce)

        return focal_loss

    def dice_plus_binary_cross_entropy_loss(self, y_true, y_predicted):

        """ compute the average of the sum of dice and binary cross entropy loss.

        Args:

            y_true: reference target.

            y_predicted:  predicted target by the model.

        Returns:

            Returns the average of the sum of dice and binary cross entropy losses.

        """

        loss = 0.5 * (self.dice_loss(y_true, y_predicted) + self.binary_cross_entropy_loss(y_true=y_true,

            y_predicted=y_predicted))

        return loss

    def dice_plus_focal_loss(self, y_true: ndarray, y_predicted: ndarray) -> float:

        """ compute the sum of the dice and focal loss

        Args:

            y_true: reference target.

            y_predicted:  predicted target by the model.

        Returns:

            Returns the sum of the dice and focal loss.

        """

        return self.dice_loss(y_true, y_predicted) + self.binary_focal_loss(y_true, y_predicted)

    @staticmethod

    def iou_loss(y_true, y_predicted, smooth=1e-8):

        """ compute the intersection over union loss.

        Args:

            y_true: reference target.

            y_predicted:  predicted target by the model.

            smooth: avoid division by zero.

        Returns:

            Returns intersection over union loss.

        """

        y_true = K.flatten(y_true)

        y_predicted = K.flatten(y_predicted)

        intersection = K.sum(K.dot(y_true, y_predicted))

        total = K.sum(y_true) + K.sum(y_predicted)

        union = total - intersection

        iou = (intersection + smooth) / (union + smooth)

        return iou