import torch
import torch.nn as nn


class DiceLoss(nn.Module):
    """Dice loss with optional pixel-wise weighting.

    Modified from segmentation_models_pytorch with added weight function
    support for dynamic pixel-wise weighting (e.g., SDF-based weighting).

    Args:
        weight_func: Optional function(pred, target) -> weights for pixel weighting.
        from_logits: If True, apply sigmoid to predictions.
        smooth: Smoothing factor to avoid division by zero.
    """

    def __init__(self, weight_func=None, from_logits=True, smooth=0.0):
        super().__init__()
        self.weight_func = weight_func
        self.smooth = smooth
        self.eps = 1e-7
        self.from_logits = from_logits

    def forward(self, pred, target, weights=None):
        """Compute the Dice loss.

        Args:
            pred: Predictions (B, C, H, W). Logits if from_logits=True.
            target: Target mask (B, C, H, W).
            weights: Optional precomputed pixel weights. If provided, overrides
                weight_func (useful when combining with other weighted losses).

        Returns:
            Dice loss value.
        """
        assert pred.size() == target.size()
        if self.from_logits:
            pred = torch.sigmoid(pred)

        if weights is not None:
            w = weights
        elif self.weight_func is not None:
            w = self.weight_func(pred, target)
        else:
            w = 1.0
        num = 2 * (w * pred * target).sum() + self.smooth
        den = ((w * pred + w * target).sum() + self.smooth).clamp(min=self.eps)
        return 1.0 - (num / den)
