import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Callable, Union


class DiceLoss(nn.Module):
    """
    Dice loss for multi-class segmentation tasks.

    Computes the soft Dice loss (1 - Dice coefficient) per batch.
    The Dice coefficient is calculated as:
        Dice = (2 * |X ∩ Y| + smooth) / (|X| + |Y| + smooth)

    The loss is averaged over all classes and batch elements.

    Args:
        smooth (float, optional): Smoothing factor to avoid division by zero.
            Defaults to 1e-5.

    Shape:
        - pred: (N, C, H, W) – raw logits for each class.
        - target: (N, H, W) – ground truth class indices (0 <= target < C).
        - Output: scalar loss value (averaged over batch and classes).

    Examples:
        >>> criterion = DiceLoss(smooth=1e-6)
        >>> logits = torch.randn(4, 10, 128, 128)   # batch of 4, 10 classes
        >>> target = torch.randint(0, 10, (4, 128, 128))
        >>> loss = criterion(logits, target)
    """

    def __init__(self, smooth: float = 1e-5) -> None:
        super().__init__()
        self.smooth = smooth

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Compute Dice loss.

        Args:
            pred (torch.Tensor): Raw logits of shape (N, C, H, W).
            target (torch.Tensor): Ground truth class indices of shape (N, H, W).

        Returns:
            torch.Tensor: Scalar loss (1 - mean Dice coefficient).
        """

        pred = torch.softmax(pred, dim=1)

        # Create one-hot encoding of target
        target_one_hot = F.one_hot(target, num_classes=pred.shape[1])  # (N, H, W, C)
        target_one_hot = target_one_hot.permute(0, 3, 1, 2).float()    # (N, C, H, W)

        # Intersection and union sums over spatial dimensions (H, W)
        intersection = (pred * target_one_hot).sum(dim=(2, 3))  # (N, C)
        union = pred.sum(dim=(2, 3)) + target_one_hot.sum(dim=(2, 3))  # (N, C)

        # Dice coefficient per class per sample
        dice = (2. * intersection + self.smooth) / (union + self.smooth)  # (N, C)

        # Average over classes and batch, then return loss
        return 1 - dice.mean()


_dice_loss = DiceLoss()


def get_loss(name: str) -> Callable:
    """
    Factory function to retrieve a loss function by name.

    Args:
        name (str): Name of the loss. One of:
            - 'ce'        : CrossEntropyLoss
            - 'dice'      : DiceLoss
            - 'combined'  : Sum of CrossEntropyLoss and DiceLoss

    Returns:
        Callable: A loss function that takes (predictions, targets) and returns a scalar loss.

    Raises:
        ValueError: If the requested loss name is not recognized.

    Examples:
        >>> criterion = get_loss('combined')
        >>> loss = criterion(logits, target)
    """
    if name == 'ce':
        return nn.CrossEntropyLoss()
    elif name == 'dice':
        # Return a new instance each time (also acceptable)
        return DiceLoss()
    elif name == 'combined':
        # Use the pre‑instantiated dice loss to avoid recreating it on every forward call
        return lambda p, t: nn.CrossEntropyLoss()(p, t) + _dice_loss(p, t)
    else:
        raise ValueError(f"Unknown loss: {name}")