"""
Optional: compute the semantic loss at a reduced resolution.

Ultralytics' SemanticSegmentationLoss upsamples the model's stride-8 logits to
the FULL mask resolution before the CE/Dice terms:

    preds = F.interpolate(preds, size=masks.shape[1:], ...)   # 64x64 -> 512x512

At batch 32 / imgsz 512 / 19 classes that materializes a 159M-element tensor,
softmaxes it, then boolean-index-gathers it (a data-dependent shape that forces
an MPS sync) -- twice, counting the aux head. Measured, the loss costs more than
the entire forward+backward of the network.

This patch computes the loss at `loss_size` instead (default 1/2 mask res):
predictions are upsampled only to loss_size and masks are nearest-downsampled to
match. Cost falls ~quadratically with loss_size.

Trade-off: less sub-grid boundary supervision, so thin classes (eyes, brows,
lips) can suffer. Computing loss below the mask resolution is standard practice
in many semseg frameworks, but it IS a quality/speed trade -- validate before
trusting it. loss_size=0 restores stock behaviour.

Usage:
    import loss_patch; loss_patch.apply(loss_size=256)
"""
import torch
import torch.nn.functional as F
from ultralytics.utils import loss as _loss

_orig_forward = _loss.SemanticSegmentationLoss.forward
_applied = False


def apply(loss_size: int = 256):
    """Patch SemanticSegmentationLoss.forward to evaluate at `loss_size`."""
    global _applied
    if loss_size <= 0:
        if _applied:
            _loss.SemanticSegmentationLoss.forward = _orig_forward
            _applied = False
        return False

    def forward(self, preds, batch):
        aux_logits = None
        if isinstance(preds, tuple):
            preds, aux_logits = preds
        masks = batch["semantic_mask"].to(preds.device)

        h, w = masks.shape[1:]
        th, tw = min(loss_size, h), min(loss_size, w)
        if (th, tw) != (h, w):
            # nearest keeps label values intact; 255 (ignore) survives too
            masks = F.interpolate(masks.float().unsqueeze(1), size=(th, tw),
                                  mode="nearest").squeeze(1).to(masks.dtype)

        valid = masks.reshape(-1) != 255
        if preds.shape[2:] != (th, tw):
            preds = F.interpolate(preds, size=(th, tw), mode="bilinear", align_corners=False)

        ce_loss = self._ce_loss(preds, masks, valid)
        dice_loss = self._dice_loss(preds, masks, valid)
        total = ce_loss + dice_loss

        aux_loss = torch.tensor(0.0, device=preds.device, dtype=ce_loss.dtype)
        if aux_logits is not None:
            if aux_logits.shape[2:] != (th, tw):
                aux_logits = F.interpolate(aux_logits, size=(th, tw), mode="bilinear",
                                           align_corners=False)
            aux_loss = self._ce_loss(aux_logits, masks, valid) * 0.4
            total += aux_loss

        return total * preds.shape[0], {"ce_loss": ce_loss.detach(),
                                        "dice_loss": dice_loss.detach(),
                                        "aux_loss": aux_loss.detach()}

    _loss.SemanticSegmentationLoss.forward = forward
    _applied = True
    return True