import torch

from .DiceLoss import DiceLoss
from .SDFWeightedBCELoss import SDFWeightedBCELoss
from .SDFWeightedDiceLoss import SDFWeightedDiceLoss
from .utils import sdf_difference_map


class DiSCoLoss(torch.nn.Module):
    """Distance-Scaled Combination (DiSCo) loss for segmentation.

    Combines base losses (BCE + Dice) with SDF-weighted losses to put heavier
    penalties on false positives/negatives that are farther from boundaries.

    The SDF computation is done once and shared between both SDF-weighted losses
    for efficiency.

    Args:
        normalize: If True, normalize SDFs to [-1, 1] range.
        baseloss_weight: Weight for base losses (BCE + Dice).
        sdfweighted_weight: Weight for SDF-weighted losses.
        clip_negatives: If True, clip negative SDF values to 0.
    """

    def __init__(
        self,
        normalize: bool = True,
        baseloss_weight: float = 1.0,
        sdfweighted_weight: float = 1.0,
        clip_negatives: bool = True,
    ):
        super().__init__()
        self.baseloss_weight = baseloss_weight
        self.sdfweighted_weight = sdfweighted_weight
        self.normalize = normalize
        self.clip_negatives = clip_negatives

        # Base losses
        self.bceloss = torch.nn.BCEWithLogitsLoss()
        self.diceloss = DiceLoss()

        # SDF-weighted losses (will receive precomputed weights)
        self.sdfwbce = SDFWeightedBCELoss(
            reduction="mean", normalize=normalize, clip_negatives=clip_negatives
        )
        self.sdfwdice = SDFWeightedDiceLoss(
            from_logits=True, normalize=normalize, clip_negatives=clip_negatives
        )

    def forward(self, pred, target):
        """Compute the DiSCo loss.

        Args:
            pred: Prediction logits (B, C, H, W).
            target: Target mask (B, C, H, W).

        Returns:
            Combined loss value.
        """
        # Base losses (no SDF computation)
        baseloss = self.bceloss(pred, target) + self.diceloss(pred, target)

        # Compute SDF weights ONCE and reuse for both SDF-weighted losses
        # This halves the SDF computation time
        sdf_weights = sdf_difference_map(
            pred > 0,  # Binarize logits
            target,
            normalize=self.normalize,
            clip_negatives=self.clip_negatives,
        )

        # SDF-weighted losses with precomputed weights
        sdf_weighted_loss = self.sdfwbce(
            pred, target, sdf_weights=sdf_weights
        ) + self.sdfwdice(pred, target, sdf_weights=sdf_weights)

        total_loss = (
            self.baseloss_weight * baseloss
            + self.sdfweighted_weight * sdf_weighted_loss
        )
        return total_loss
