"""Custom loss functions for coronary segmentation with proper MONAI interface."""

import torch
import torch.nn as nn
from monai.losses.cldice import soft_dice, soft_skel
from monai.networks.utils import one_hot


class DiceclDiceLoss(nn.Module):
    """
    Combined Dice and clDice loss with proper MONAI interface.

    This wrapper addresses the interface limitations of MONAI's SoftDiceclDiceLoss
    (see https://github.com/Project-MONAI/MONAI/issues/8239).

    Uses (input, target) argument order compatible with MONAI SupervisedTrainer.
    Handles softmax and one-hot encoding internally.

    Args:
        iter_: Number of iterations for soft skeleton computation.
        alpha: Weight for clDice loss (1-alpha for Dice loss).
        smooth: Smoothing parameter to avoid division by zero.
        include_background: Whether to include background class in loss computation.
        to_onehot_y: Whether to convert target to one-hot encoding.
        softmax: Whether to apply softmax to input predictions.
    """

    def __init__(
        self,
        iter_: int = 50,
        alpha: float = 0.5,
        smooth: float = 1.0,
        include_background: bool = False,
        to_onehot_y: bool = True,
        softmax: bool = True,
    ):
        super().__init__()
        self.iter = iter_
        self.alpha = alpha
        self.smooth = smooth
        self.include_background = include_background
        self.to_onehot_y = to_onehot_y
        self.softmax = softmax

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        n_classes = input.shape[1]

        if self.softmax:
            y_pred = torch.softmax(input, dim=1)
        else:
            y_pred = input

        if self.to_onehot_y:
            y_true = one_hot(target, num_classes=n_classes)
        else:
            y_true = target

        if not self.include_background:
            y_pred = y_pred[:, 1:, ...]
            y_true = y_true[:, 1:, ...]

        dice_loss = 1.0 - soft_dice(y_true, y_pred, self.smooth)

        skel_pred = soft_skel(y_pred, self.iter)
        skel_true = soft_skel(y_true, self.iter)

        tprec = (torch.sum(skel_pred * y_true) + self.smooth) / (
            torch.sum(skel_pred) + self.smooth
        )
        tsens = (torch.sum(skel_true * y_pred) + self.smooth) / (
            torch.sum(skel_true) + self.smooth
        )
        cl_dice_loss = 1.0 - 2.0 * (tprec * tsens) / (tprec + tsens)

        return (1.0 - self.alpha) * dice_loss + self.alpha * cl_dice_loss