from functools import partial
from .DiceLoss import DiceLoss
from .utils import sdf_difference_map


class SDFWeightedDiceLoss(DiceLoss):
    """Dice loss weighted by SDF difference between prediction and target.

    The loss applies heavier penalties to false positives/negatives that are
    farther from the correct boundary.

    Args:
        from_logits: If True, apply sigmoid to predictions.
        normalize: If True, normalize SDFs to [-1, 1] range.
        clip_negatives: If True, clip negative SDF values to 0.
    """

    def __init__(
        self,
        from_logits: bool = True,
        normalize: bool = True,
        clip_negatives: bool = False,
    ):
        super().__init__(
            weight_func=partial(
                sdf_difference_map, normalize=normalize, clip_negatives=clip_negatives
            ),
            from_logits=from_logits,
        )
        self.normalize = normalize
        self.clip_negatives = clip_negatives

    def forward(self, pred, target, sdf_weights=None):
        """Compute the SDF-weighted Dice loss.

        Args:
            pred: Predictions (B, C, H, W). Logits if from_logits=True.
            target: Target mask (B, C, H, W).
            sdf_weights: Optional precomputed SDF weights. If provided, skips
                SDF computation (useful when combining with other SDF losses).

        Returns:
            Weighted Dice loss.
        """
        return super().forward(pred, target, weights=sdf_weights)
