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) #Applies softmax along the class dimension (dim=1)
        target_one_hot = F.one_hot(target,num_classes=pred.shape[1])  # shape: (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()

    def get_loss(self,name):
        if name == 'ce':
            return nn.CrossEntropyLoss()
        elif name == 'dice':
            return DiceLoss()
        elif name == 'combined':
            return lambda p, t: nn.CrossEntropyLoss()(p, t) + DiceLoss()(p, t)
        else:
            raise ValueError(f"Unknown loss: {name}")