"""损失函数。"""

from __future__ import annotations

import torch
from torch import nn
from torch.nn import functional as F

from utils.labels import MULTILABEL, check_label_mode


def gaussian_focal_loss(
    pred_logits: torch.Tensor,
    gt_heatmap: torch.Tensor,
    alpha: float = 2.0,
    beta: float = 4.0,
    eps: float = 1e-4,
) -> torch.Tensor:
    """CenterNet 的 penalty-reduced focal loss，按正样本数归一化。"""
    pred = pred_logits.sigmoid().clamp(eps, 1 - eps)
    pos_mask = gt_heatmap.eq(1.0).float()
    neg_mask = 1.0 - pos_mask

    pos_loss = -((1 - pred) ** alpha) * torch.log(pred) * pos_mask
    neg_loss = -((1 - gt_heatmap) ** beta) * (pred**alpha) * torch.log(1 - pred) * neg_mask

    num_pos = pos_mask.sum().clamp(min=1.0)
    return (pos_loss.sum() + neg_loss.sum()) / num_pos


def masked_l1_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """只在中心点位置计算 L1。mask: (B, H, W)。"""
    mask = mask.unsqueeze(1).expand_as(pred)
    loss = F.l1_loss(pred * mask, target * mask, reduction="sum")
    return loss / mask.sum().clamp(min=1.0)


class DetectionLoss(nn.Module):
    """heatmap + wh + offset (+ 图像级分类) 的加权和。

    分类分支按 `label_mode` 切换：
        multilabel（多类）—— 多热标签，逐类 BCE；
        multiclass（单类）—— 互斥下标，softmax 交叉熵。
    """

    def __init__(
        self,
        heatmap_weight: float = 1.0,
        wh_weight: float = 0.1,
        offset_weight: float = 1.0,
        cls_weight: float = 1.0,
        label_mode: str = MULTILABEL,
    ) -> None:
        super().__init__()
        self.label_mode = check_label_mode(label_mode)
        self.heatmap_weight = heatmap_weight
        self.wh_weight = wh_weight
        self.offset_weight = offset_weight
        self.cls_weight = cls_weight

    def forward(self, outputs: dict, batch: dict) -> dict[str, torch.Tensor]:
        target = batch["target"]
        losses = {
            "heatmap": self.heatmap_weight * gaussian_focal_loss(outputs["heatmap"], target["heatmap"]),
            "wh": self.wh_weight * masked_l1_loss(outputs["wh"], target["wh"], target["mask"]),
            "offset": self.offset_weight
            * masked_l1_loss(outputs["offset"], target["offset"], target["mask"]),
        }
        if "logits" in outputs and self.cls_weight > 0:
            losses["cls"] = self.cls_weight * self._classification_loss(
                outputs["logits"], batch["image_labels"]
            )
        losses["total"] = sum(losses.values())
        return losses

    def _classification_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        if self.label_mode == MULTILABEL:
            return F.binary_cross_entropy_with_logits(logits, targets.float())
        return F.cross_entropy(logits, targets.long())