"""指标：分类与定位。"""

from __future__ import annotations

import torch


class AverageMeter:
    def __init__(self) -> None:
        self.total = 0.0
        self.count = 0

    def update(self, value: float, n: int = 1) -> None:
        self.total += float(value) * n
        self.count += n

    @property
    def avg(self) -> float:
        return self.total / self.count if self.count else 0.0


def box_iou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) -> torch.Tensor:
    """(N, 4) 与 (M, 4) 的两两 IoU，返回 (N, M)。"""
    if boxes_a.numel() == 0 or boxes_b.numel() == 0:
        return boxes_a.new_zeros((boxes_a.shape[0], boxes_b.shape[0]))
    area_a = (boxes_a[:, 2] - boxes_a[:, 0]).clamp(min=0) * (boxes_a[:, 3] - boxes_a[:, 1]).clamp(min=0)
    area_b = (boxes_b[:, 2] - boxes_b[:, 0]).clamp(min=0) * (boxes_b[:, 3] - boxes_b[:, 1]).clamp(min=0)

    lt = torch.max(boxes_a[:, None, :2], boxes_b[None, :, :2])
    rb = torch.min(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
    wh = (rb - lt).clamp(min=0)
    inter = wh[..., 0] * wh[..., 1]
    return inter / (area_a[:, None] + area_b[None, :] - inter).clamp(min=1e-6)


class DetectionMetric:
    """按 IoU 阈值做贪心匹配，统计 precision / recall / F1 与平均 IoU。

    这是工程上够用的近似指标；正式对比模型时建议再跑一次 COCO mAP。
    """

    def __init__(self, iou_threshold: float = 0.5) -> None:
        self.iou_threshold = iou_threshold
        self.reset()

    def reset(self) -> None:
        self.tp = 0
        self.fp = 0
        self.fn = 0
        self.iou_sum = 0.0

    def update(
        self,
        pred_boxes: torch.Tensor,
        pred_labels: torch.Tensor,
        gt_boxes: torch.Tensor,
        gt_labels: torch.Tensor,
    ) -> None:
        if gt_boxes.numel() == 0:
            self.fp += int(pred_boxes.shape[0])
            return
        if pred_boxes.numel() == 0:
            self.fn += int(gt_boxes.shape[0])
            return

        ious = box_iou(pred_boxes.float().cpu(), gt_boxes.float().cpu())
        same_class = pred_labels.cpu()[:, None] == gt_labels.cpu()[None, :]
        ious = ious * same_class.float()

        matched_gt: set[int] = set()
        for pred_idx in range(ious.shape[0]):
            best_iou, best_gt = ious[pred_idx].max(dim=0)
            gt_idx = int(best_gt)
            if float(best_iou) >= self.iou_threshold and gt_idx not in matched_gt:
                matched_gt.add(gt_idx)
                self.tp += 1
                self.iou_sum += float(best_iou)
            else:
                self.fp += 1
        self.fn += int(gt_boxes.shape[0]) - len(matched_gt)

    def compute(self) -> dict[str, float]:
        precision = self.tp / max(self.tp + self.fp, 1)
        recall = self.tp / max(self.tp + self.fn, 1)
        f1 = 2 * precision * recall / max(precision + recall, 1e-9)
        return {
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "mean_iou": self.iou_sum / max(self.tp, 1),
        }


class MultiLabelMetric:
    """多类模式：逐类阈值化后统计 micro precision / recall / F1 与整图完全正确率。"""

    def __init__(self, threshold: float = 0.5) -> None:
        self.threshold = threshold
        self.reset()

    def reset(self) -> None:
        self.tp = 0
        self.fp = 0
        self.fn = 0
        self.exact = 0
        self.total = 0

    def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
        pred = (logits.sigmoid() >= self.threshold).float().cpu()
        gt = targets.float().cpu()
        self.tp += int((pred * gt).sum())
        self.fp += int((pred * (1 - gt)).sum())
        self.fn += int(((1 - pred) * gt).sum())
        self.exact += int((pred == gt).all(dim=1).sum())
        self.total += int(gt.shape[0])

    def compute(self) -> dict[str, float]:
        precision = self.tp / max(self.tp + self.fp, 1)
        recall = self.tp / max(self.tp + self.fn, 1)
        return {
            "cls_precision": precision,
            "cls_recall": recall,
            "cls_f1": 2 * precision * recall / max(precision + recall, 1e-9),
            "cls_exact_match": self.exact / max(self.total, 1),
        }


class SingleLabelMetric:
    """单类模式：互斥分类的准确率与 macro precision / recall / F1。

    指标名与 `MultiLabelMetric` 保持一致（cls_precision / cls_recall / cls_f1），
    因此切换 label_mode 后 `train.monitor` 不需要改。
    """

    def __init__(self, num_outputs: int) -> None:
        self.num_outputs = num_outputs
        self.reset()

    def reset(self) -> None:
        self.tp = torch.zeros(self.num_outputs)
        self.fp = torch.zeros(self.num_outputs)
        self.fn = torch.zeros(self.num_outputs)
        self.correct = 0
        self.total = 0

    def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
        pred = logits.argmax(dim=1).cpu()
        gt = targets.long().cpu()
        for cls in range(self.num_outputs):
            pred_cls = pred == cls
            gt_cls = gt == cls
            self.tp[cls] += int((pred_cls & gt_cls).sum())
            self.fp[cls] += int((pred_cls & ~gt_cls).sum())
            self.fn[cls] += int((~pred_cls & gt_cls).sum())
        self.correct += int((pred == gt).sum())
        self.total += int(gt.numel())

    def compute(self) -> dict[str, float]:
        # 只对验证集中出现过的类别做 macro 平均，避免缺席类别把指标拉到 0。
        present = (self.tp + self.fn) > 0
        precision = self.tp / (self.tp + self.fp).clamp(min=1)
        recall = self.tp / (self.tp + self.fn).clamp(min=1)
        f1 = 2 * precision * recall / (precision + recall).clamp(min=1e-9)
        if present.any():
            precision, recall, f1 = precision[present], recall[present], f1[present]
        return {
            "cls_accuracy": self.correct / max(self.total, 1),
            "cls_precision": float(precision.mean()),
            "cls_recall": float(recall.mean()),
            "cls_f1": float(f1.mean()),
        }