"""分类任务的训练、评估与指标。

数据极不均衡（normal 2679 张 vs others 65 张，差 41 倍），所以：
- 主指标用 **macro-F1** 而不是准确率——全预测成 normal 也能拿到 24% 准确率，但 macro-F1 会很低；
- 损失支持按类别频率加权，避免小类被淹没。
"""

from __future__ import annotations

import time
import warnings
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader

from engine.checkpoint import load_checkpoint, save_checkpoint
from engine.metrics import AverageMeter
from engine.optim import build_optimizer, build_scheduler
from engine.visualizer import SwanLabLogger
from utils.device import PeakMemory, resolve_amp
from utils.logging import get_logger

logger = get_logger("engine.classification")


class ClassificationMetric:
    """混淆矩阵 + 由它导出的各项指标。

    准确率在不均衡数据上会骗人，所以同时给 macro-F1（各类等权）与
    weighted-F1（按样本数加权），并保留逐类明细供人工判断哪些类没学会。
    """

    def __init__(self, classes: list[str]) -> None:
        self.classes = list(classes)
        self.reset()

    def reset(self) -> None:
        n = len(self.classes)
        self.confusion = torch.zeros(n, n, dtype=torch.long)  # 行=真实，列=预测

    def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
        preds = logits.argmax(dim=1).cpu()
        for target, pred in zip(targets.cpu().tolist(), preds.tolist(), strict=True):
            self.confusion[target, pred] += 1

    def per_class(self) -> dict[str, dict[str, float]]:
        report = {}
        for index, name in enumerate(self.classes):
            tp = int(self.confusion[index, index])
            support = int(self.confusion[index].sum())
            predicted = int(self.confusion[:, index].sum())
            precision = tp / predicted if predicted else 0.0
            recall = tp / support if support else 0.0
            f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
            report[name] = {
                "precision": precision,
                "recall": recall,
                "f1": f1,
                "support": support,
            }
        return report

    def compute(self) -> dict[str, float]:
        total = int(self.confusion.sum())
        if not total:
            return {"accuracy": 0.0, "macro_f1": 0.0, "weighted_f1": 0.0, "balanced_accuracy": 0.0}
        report = self.per_class()
        f1s = [r["f1"] for r in report.values()]
        recalls = [r["recall"] for r in report.values() if r["support"]]
        supports = [r["support"] for r in report.values()]
        weighted = sum(r["f1"] * r["support"] for r in report.values()) / total
        return {
            "accuracy": float(self.confusion.diagonal().sum()) / total,
            "macro_f1": sum(f1s) / len(f1s),
            "weighted_f1": weighted,
            # 各类召回的均值：小类分错会明显拉低它，比准确率更能反映"每类都学会了吗"
            "balanced_accuracy": sum(recalls) / len(recalls) if recalls else 0.0,
            "num_classes_with_support": sum(1 for s in supports if s),
        }


def class_weights(counts: list[int], mode: str | None = "balanced") -> torch.Tensor | None:
    """按类别频率算损失权重。

    `balanced` 用逆频率并归一化到均值 1：小类的每个样本被放大，
    否则 normal（占 24%）会主导梯度，小类学不动。
    """
    if not mode or str(mode).lower() in ("none", "false"):
        return None
    if str(mode).lower() != "balanced":
        raise ValueError(f"未知 class_weight: {mode}，可选 balanced / none")
    counts_tensor = torch.tensor(counts, dtype=torch.float32).clamp(min=1)
    weights = counts_tensor.sum() / (len(counts) * counts_tensor)
    return weights / weights.mean()


class ClassificationTrainer:
    """Args:
    model: 已 `.to(device)` 的分类模型。
    train_loader / val_loader: 数据加载器。
    cfg: `train` 段配置（epochs / optimizer / scheduler / amp / class_weight / monitor 等）。
    classes: 类别名，顺序必须与标签下标一致。
    """

    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader | None,
        cfg: Mapping[str, Any],
        device: torch.device,
        work_dir: str | Path,
        classes: list[str],
        train_counts: list[int] | None = None,
        full_config: dict | None = None,
    ) -> None:
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.cfg = dict(cfg)
        self.device = device
        self.classes = list(classes)
        self.work_dir = Path(work_dir)
        self.work_dir.mkdir(parents=True, exist_ok=True)
        self.full_config = full_config

        self.epochs = int(self.cfg.get("epochs", 30))
        self.log_interval = int(self.cfg.get("log_interval", 20))
        self.clip_grad = float(self.cfg.get("clip_grad", 0.0))
        self.monitor = str(self.cfg.get("monitor", "macro_f1"))
        self.label_smoothing = float(self.cfg.get("label_smoothing", 0.0))

        weights = class_weights(train_counts or [], self.cfg.get("class_weight", "balanced"))
        self.class_weight = weights.to(device) if weights is not None else None
        if self.class_weight is not None:
            extremes = sorted(zip(self.classes, weights.tolist(), strict=True), key=lambda kv: -kv[1])
            logger.info(
                "类别权重（逆频率）：最大 %s=%.2f，最小 %s=%.2f",
                extremes[0][0], extremes[0][1], extremes[-1][0], extremes[-1][1],
            )

        self.optimizer = build_optimizer(model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3}))
        self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
        self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
        self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
        self.peak_memory = PeakMemory(device)

        self.metric = ClassificationMetric(self.classes)
        self.best_metric = float("-inf")
        self.start_epoch = 0
        self.global_step = 0
        self.visualizer = SwanLabLogger(
            work_dir=self.work_dir, config=full_config, **dict(self.cfg.get("swanlab", {}))
        )

    def _loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        return F.cross_entropy(
            logits, targets, weight=self.class_weight, label_smoothing=self.label_smoothing
        )

    def train_one_epoch(self, epoch: int) -> dict[str, float]:
        self.model.train()
        loss_meter, acc_meter = AverageMeter(), AverageMeter()
        start = time.time()

        for step, batch in enumerate(self.train_loader):
            images = batch["images"].to(self.device, non_blocking=True)
            targets = batch["labels"].to(self.device, non_blocking=True)

            with torch.amp.autocast(self.device.type, enabled=self.amp):
                logits = self.model(images)
                loss = self._loss(logits, targets)

            self.optimizer.zero_grad(set_to_none=True)
            self.scaler.scale(loss).backward()
            if self.clip_grad > 0:
                self.scaler.unscale_(self.optimizer)
                nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_grad)
            self.scaler.step(self.optimizer)
            self.scaler.update()

            self.peak_memory.update()
            batch_size = images.shape[0]
            loss_meter.update(float(loss.detach()), batch_size)
            acc_meter.update(float((logits.argmax(1) == targets).float().mean()), batch_size)
            self.global_step += 1

            if step % self.log_interval == 0:
                logger.info(
                    "epoch %d [%d/%d] loss=%.4f acc=%.4f lr=%.2e",
                    epoch, step, len(self.train_loader), loss_meter.avg, acc_meter.avg,
                    self.optimizer.param_groups[0]["lr"],
                )
                self.visualizer.log(
                    {"loss": loss_meter.avg, "acc": acc_meter.avg,
                     "lr": self.optimizer.param_groups[0]["lr"]},
                    step=self.global_step, prefix="train/",
                )

        logger.info("epoch %d 训练耗时 %.1fs", epoch, time.time() - start)
        return {"loss": loss_meter.avg, "acc": acc_meter.avg}

    @torch.no_grad()
    def evaluate(self, loader: DataLoader | None = None) -> dict[str, float]:
        loader = loader or self.val_loader
        if loader is None:
            return {}
        self.model.eval()
        self.metric.reset()
        loss_meter = AverageMeter()

        for batch in loader:
            images = batch["images"].to(self.device, non_blocking=True)
            targets = batch["labels"].to(self.device, non_blocking=True)
            logits = self.model(images)
            loss_meter.update(float(self._loss(logits, targets)), images.shape[0])
            self.metric.update(logits, targets)

        metrics = self.metric.compute()
        metrics["loss"] = loss_meter.avg
        return metrics

    def fit(self) -> float:
        logger.info("开始训练：共 %d 个 epoch，%d 类，输出目录 %s",
                    self.epochs, len(self.classes), self.work_dir)
        for epoch in range(self.start_epoch, self.epochs):
            self.train_one_epoch(epoch)
            if self.scheduler is not None:
                self.scheduler.step()

            metrics = self.evaluate()
            if metrics:
                logger.info(
                    "epoch %d 验证: %s",
                    epoch,
                    " ".join(f"{k}={v:.4f}" for k, v in metrics.items() if isinstance(v, float)),
                )
                self.visualizer.log(metrics, step=self.global_step, prefix="val/")

            self._save("last.pth", epoch)
            score = metrics.get(self.monitor)
            if score is not None and score > self.best_metric:
                self.best_metric = score
                self._save("best.pth", epoch)
                logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score)

        self.visualizer.finish()
        return self.best_metric

    def _save(self, name: str, epoch: int) -> None:
        save_checkpoint(
            self.work_dir / name, self.model, self.optimizer, epoch=epoch,
            best_metric=self.best_metric, config=self.full_config,
            scheduler=self.scheduler, global_step=self.global_step, train_cfg=dict(self.cfg),
        )

    def resume(self, path: str | Path) -> dict:
        """续训。与检测训练同样的语义：检查点只提供进度，超参以配置为准。"""
        state = load_checkpoint(path, self.model, map_location=self.device)
        if state.get("optimizer"):
            intended = [
                {k: v for k, v in group.items() if k != "params"}
                for group in self.optimizer.param_groups
            ]
            self.optimizer.load_state_dict(state["optimizer"])
            for group, wanted in zip(self.optimizer.param_groups, intended, strict=True):
                group.update(wanted)
        self.start_epoch = int(state.get("epoch", -1)) + 1
        self.global_step = int(state.get("global_step", 0))
        best = state.get("best_metric")
        self.best_metric = float(best) if best is not None else float("-inf")

        if self.scheduler is not None:
            for group in self.optimizer.param_groups:
                group["lr"] = group.get("initial_lr", group["lr"])
            self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                for _ in range(self.start_epoch):
                    self.scheduler.step()
        logger.info("已从 %s 恢复：起始 epoch=%d，best_%s=%.4f",
                    path, self.start_epoch, self.monitor, self.best_metric)
        return state