import torch
from typing import Literal
from .utils import sdf_difference_weighted_bceloss_from_logits


class SDFWeightedBCELoss(torch.nn.Module):
    """Compute BCE 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:
        reduction: Reduction mode ('none', 'mean', 'sum').
        normalize: If True, normalize SDFs to [-1, 1] range.
        clip_negatives: If True, clip negative SDF values to 0.
    """

    def __init__(
        self,
        reduction: Literal["mean", "sum", "none"] = "mean",
        normalize: bool = True,
        clip_negatives: bool = False,
    ):
        super().__init__()
        self.reduction = reduction
        self.normalize = normalize
        self.clip_negatives = clip_negatives

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

        Args:
            pred: Prediction logits (B, C, H, W).
            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 BCE loss.
        """
        return sdf_difference_weighted_bceloss_from_logits(
            pred,
            target,
            reduction=self.reduction,
            normalize=self.normalize,
            clip_negatives=self.clip_negatives,
            sdf_weights=sdf_weights,
        )
