import torch
import torch.nn as nn
import torch.nn.functional as F

class DiceLoss(nn.Module):
    def __init__(self, smooth=1e-5):
        super().__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        # pred: logits (B, C, H, W)
        # target: (B, H, W) with class indices
        pred = torch.softmax(pred, dim=1)
        target_one_hot = F.one_hot(target, num_classes=pred.shape[1])  # (B, H, W, C)
        target_one_hot = target_one_hot.permute(0, 3, 1, 2).float()

        intersection = (pred * target_one_hot).sum(dim=(2, 3))
        union = pred.sum(dim=(2, 3)) + target_one_hot.sum(dim=(2, 3))
        dice = (2. * intersection + self.smooth) / (union + self.smooth)
        return 1 - dice.mean()

# Instantiate a single DiceLoss instance for reuse
_dice_loss = DiceLoss()

def get_loss(name):
    if name == 'ce':
        return nn.CrossEntropyLoss()
    elif name == 'dice':
        return DiceLoss()          # returns a new instance (also fine)
    elif name == 'combined':
        # Reuse the pre‑instantiated dice loss to avoid recreating every forward
        return lambda p, t: nn.CrossEntropyLoss()(p, t) + _dice_loss(p, t)
    else:
        raise ValueError(f"Unknown loss: {name}")