"""评估：跑一遍验证集，输出分类与定位指标。"""

from __future__ import annotations

from typing import Any

import torch
from torch import nn
from torch.utils.data import DataLoader

from engine.metrics import AverageMeter, DetectionMetric, MultiLabelMetric, SingleLabelMetric
from models.postprocess import decode_heatmap
from utils.labels import MULTILABEL
from utils.logging import get_logger

logger = get_logger("engine.evaluator")


class Evaluator:
    def __init__(
        self,
        model: nn.Module,
        dataloader: DataLoader,
        device: torch.device,
        criterion: nn.Module | None = None,
        score_threshold: float = 0.3,
        iou_threshold: float = 0.5,
        topk: int = 50,
    ) -> None:
        self.model = model
        self.dataloader = dataloader
        self.device = device
        self.criterion = criterion
        self.score_threshold = score_threshold
        self.topk = topk
        self.det_metric = DetectionMetric(iou_threshold)
        # 分类指标按模型的 label_mode 选择：多类用多标签指标，单类用互斥分类指标。
        self.label_mode = getattr(model, "label_mode", MULTILABEL)
        if self.label_mode == MULTILABEL:
            self.cls_metric: MultiLabelMetric | SingleLabelMetric = MultiLabelMetric()
        else:
            num_outputs = model.cls_head.num_outputs if model.cls_head is not None else 1
            self.cls_metric = SingleLabelMetric(num_outputs)

    @torch.no_grad()
    def evaluate(self) -> dict[str, float]:
        self.model.eval()
        self.det_metric.reset()
        self.cls_metric.reset()
        loss_meter = AverageMeter()

        stride = getattr(self.model, "out_stride", 4)
        for batch in self.dataloader:
            images = batch["images"].to(self.device, non_blocking=True)
            outputs = self.model(images)

            if self.criterion is not None:
                batch_on_device = _to_device(batch, self.device)
                losses = self.criterion(outputs, batch_on_device)
                loss_meter.update(float(losses["total"]), images.shape[0])

            detections = decode_heatmap(
                outputs["heatmap"],
                outputs["wh"],
                outputs["offset"],
                stride=stride,
                topk=self.topk,
                score_threshold=self.score_threshold,
            )
            for i, det in enumerate(detections):
                self.det_metric.update(det["boxes"], det["labels"], batch["boxes"][i], batch["labels"][i])

            if "logits" in outputs:
                self.cls_metric.update(outputs["logits"], batch["image_labels"])

        metrics: dict[str, float] = {**self.det_metric.compute(), **self.cls_metric.compute()}
        if self.criterion is not None:
            metrics["loss"] = loss_meter.avg
        return metrics


def _to_device(batch: dict, device: torch.device) -> dict:
    moved: dict[str, Any] = dict(batch)
    moved["images"] = batch["images"].to(device, non_blocking=True)
    moved["target"] = {k: v.to(device, non_blocking=True) for k, v in batch["target"].items()}
    moved["image_labels"] = batch["image_labels"].to(device, non_blocking=True)
    return moved