# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_vision_loss_functions.ipynb.

# %% auto #0
__all__ = ['CustomLoss', 'TverskyFocalLoss']

# %% ../nbs/04_vision_loss_functions.ipynb #e8eee6b3-75c3-4468-8645-8f6a42a76bfe
import torch
from .vision_core import *
from .vision_data import pred_to_binary_mask, batch_pred_to_multiclass_mask
from monai.losses import TverskyLoss, FocalLoss
from torch.nn.modules.loss import _Loss

# %% ../nbs/04_vision_loss_functions.ipynb #e0c0c220-aaeb-46c6-8d18-f72cd9da0555
class CustomLoss:
    """A custom loss wrapper class for loss functions to allow them to work with
    the 'show_results' method in fastai. 
    """

    def __init__(self, loss_func):
        """Constructs CustomLoss object."""
        
        self.loss_func = loss_func

    def __call__(self, pred, targ):
        """Computes the loss for given predictions and targets."""
        
        if isinstance(pred, MedBase):
            pred, targ = torch.Tensor(pred.cpu()), torch.Tensor(targ.cpu().float())
            
        return self.loss_func(pred, targ)

    def activation(self, x):
        return x
    
    def decodes(self, x) -> torch.Tensor:
        """Converts model output to target format.
        
        Args:
            x: Activations for each class with dimensions [B, C, W, H, D].

        Returns:
            The predicted mask.
        """
        
        n_classes = x.shape[1]
        if n_classes == 1: 
            x = pred_to_binary_mask(x)
        else: 
            x,_ = batch_pred_to_multiclass_mask(x)

        return x

# %% ../nbs/04_vision_loss_functions.ipynb #5052c7bc-3d9a-4e34-8b64-bceaf2fdc7b6
class TverskyFocalLoss(_Loss):
    """
    Compute Tversky loss with a focus parameter, gamma, applied.
    The details of Tversky loss is shown in ``monai.losses.TverskyLoss``.
    """

    def __init__(
        self,
        include_background: bool = True,
        to_onehot_y: bool = False,
        sigmoid: bool = False,
        softmax: bool = False,
        gamma: float = 2,
        alpha: float = 0.5, 
        beta: float = 0.99):
        """
        Args:
            include_background: if to calculate loss for the background class.
            to_onehot_y: whether to convert `y` into one-hot format.
            sigmoid: if True, apply a sigmoid function to the prediction.
            softmax: if True, apply a softmax function to the prediction.
            gamma: the focal parameter, it modulates the loss with regards to 
                how far the prediction is from target.
            alpha: the weight of false positive in Tversky loss calculation.
            beta: the weight of false negative in Tversky loss calculation.
        """
            
        super().__init__()
        self.tversky = TverskyLoss(
            to_onehot_y=to_onehot_y, 
            include_background=include_background, 
            sigmoid=sigmoid, 
            softmax=softmax, 
            alpha=alpha, 
            beta=beta
        )
        self.gamma = gamma

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Args:
            input: the shape should be [B, C, W, H, D]. The input should be the original logits.
            target: the shape should be[B, C, W, H, D].

        Raises:
            ValueError: When number of dimensions for input and target are different.
        """
        if len(input.shape) != len(target.shape):
            raise ValueError("The number of dimensions for input and target should be the same.")

        tversky_loss = self.tversky(input, target)
        total_loss: torch.Tensor = 1 - ((1 - tversky_loss)**self.gamma)

        return total_loss
