"""训练循环。支持 cuda / mps / cpu 三种后端。"""

from __future__ import annotations

import math
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.utils.data import DataLoader

from engine.checkpoint import load_checkpoint, save_checkpoint
from engine.evaluator import Evaluator, _to_device
from engine.losses import DetectionLoss
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.labels import MULTILABEL
from utils.logging import get_logger

logger = get_logger("engine.trainer")


class Trainer:
    """Args:
    model: 已经 `.to(device)` 的模型。
    train_loader / val_loader: 数据加载器，val 可为 None。
    cfg: `train` 段配置（epochs / optimizer / scheduler / amp / ...）。
    work_dir: 日志与权重输出目录。
    """

    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader | None,
        cfg: Mapping[str, Any],
        device: torch.device,
        work_dir: str | Path,
        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.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", 50))
        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", "f1"))

        # 分类损失形式跟随模型的 label_mode（yaml 里的 label_mode 决定多类 / 单类）。
        loss_cfg = dict(self.cfg.get("loss", {}))
        loss_cfg.setdefault("label_mode", getattr(model, "label_mode", MULTILABEL))
        self.criterion = DetectionLoss(**loss_cfg)
        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)

        # amp: auto 时随后端（CUDA 开、MPS 关——MPS 上 autocast 实测比 fp32 慢）
        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.evaluator = (
            Evaluator(
                model,
                val_loader,
                device,
                criterion=self.criterion,
                **dict(self.cfg.get("eval", {})),
            )
            if val_loader is not None
            else None
        )
        self.visualizer = SwanLabLogger(
            work_dir=self.work_dir,
            config=full_config,
            **dict(self.cfg.get("swanlab", {})),
        )
        self.best_metric = float("-inf")
        self.start_epoch = 0
        self.global_step = 0

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

        for step, batch in enumerate(self.train_loader):
            batch = _to_device(batch, self.device)

            with torch.amp.autocast(self.device.type, enabled=self.amp):
                outputs = self.model(batch["images"])
                losses = self.criterion(outputs, batch)

            self.optimizer.zero_grad(set_to_none=True)
            self.scaler.scale(losses["total"]).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()

            batch_size = batch["images"].shape[0]
            step_losses = {key: float(value.detach()) for key, value in losses.items()}
            for key, value in step_losses.items():
                meters.setdefault(key, AverageMeter()).update(value, batch_size)

            self.global_step += 1
            if step % self.log_interval == 0:
                self.visualizer.log(
                    {**step_losses, "lr": self.optimizer.param_groups[0]["lr"]},
                    step=self.global_step,
                    prefix="train/",
                )
                summary = " ".join(f"{k}={m.avg:.4f}" for k, m in meters.items())
                logger.info(
                    "epoch %d [%d/%d] lr=%.2e %s",
                    epoch,
                    step,
                    len(self.train_loader),
                    self.optimizer.param_groups[0]["lr"],
                    summary,
                )

        logger.info("epoch %d 训练耗时 %.1fs", epoch, time.time() - start)
        return {key: meter.avg for key, meter in meters.items()}

    def _gpu_memory(self) -> dict[str, float]:
        """本轮训练的显存峰值（GB）；推理端的详细分析见 scripts/profile_inference.py。

        CUDA 读内置峰值计数器；MPS 没有峰值 API，用逐 step 采样的最大值（是下界）；
        CPU 不统计。
        """
        if self.device.type == "cpu":
            return {}
        stats = {
            "gpu_mem_allocated_gb": self.peak_memory.peak_mb / 1024,
            "gpu_mem_reserved_gb": self.peak_memory.reserved_mb / 1024,
        }
        self.peak_memory.reset()
        return stats

    def resume(self, path: str | Path) -> dict:
        """从检查点续训。

        **配置文件是唯一真相源**：检查点只提供"进度"——模型权重、优化器动量缓冲、
        已训到第几个 epoch、global_step、best_metric；而 lr、weight_decay、总轮数、
        调度器一律以当前 yaml 为准。否则改了配置再续训，改动会被检查点静默覆盖。
        """
        state = load_checkpoint(path, self.model, map_location=self.device)

        if state.get("optimizer"):
            # optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay
            # 一并还原成存档时的值。先抄下配置构造出的超参，加载完按组写回，
            # 只留下动量缓冲。按组写回也保住了 BN 与 bias 那组 weight_decay=0 的结构。
            intended = [
                {key: value for key, value in group.items() if key != "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 and math.isfinite(float(best)) else float("-inf")

        self._restore_scheduler()
        self._log_config_changes(state)
        logger.info(
            "已从 %s 恢复：起始 epoch=%d，best_%s=%.4f",
            path,
            self.start_epoch,
            self.monitor,
            self.best_metric,
        )
        return state

    def _restore_scheduler(self) -> None:
        """按**当前配置**重建学习率调度，再快进到已训到的 epoch。

        不恢复检查点里的调度器状态：它带着当时的 T_max 与 base_lrs，一旦 yaml 里改了
        总轮数或学习率就会被静默覆盖。cosine / step / multistep 都是 epoch 的确定性函数，
        快进与恢复状态等价，所以统一走重建这一条路。
        """
        if self.scheduler is None:
            return
        # cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr，
        # 否则从检查点带来的 lr 出发会算出一条错的曲线。
        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")  # 快进阶段的 "step before optimizer.step()" 属预期
            for _ in range(self.start_epoch):
                self.scheduler.step()

    def _log_config_changes(self, state: dict) -> None:
        """检查点里的配置与当前配置不一致时说清楚用的是哪一份。"""
        archived = state.get("train_cfg") or {}
        changes = []
        if archived.get("epochs") not in (None, self.epochs):
            changes.append(f"总轮数 {archived['epochs']} -> {self.epochs}")
        for key, value in (archived.get("optimizer") or {}).items():
            current = dict(self.cfg.get("optimizer", {})).get(key)
            if current is not None and current != value:
                changes.append(f"optimizer.{key} {value} -> {current}")
        archived_scheduler = (archived.get("scheduler") or {}).get("name")
        current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
        if archived_scheduler and archived_scheduler != current_scheduler:
            changes.append(f"scheduler {archived_scheduler} -> {current_scheduler}")

        if changes:
            logger.warning("配置与检查点不一致，按**当前配置**执行：%s", "；".join(changes))
        logger.info(
            "续训超参：lr=%.2e（initial_lr=%.2e）weight_decay=%s，共 %d 轮",
            self.optimizer.param_groups[0]["lr"],
            self.optimizer.param_groups[0].get("initial_lr", float("nan")),
            [g["weight_decay"] for g in self.optimizer.param_groups],
            self.epochs,
        )

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

            metrics: dict[str, float] = {}
            if self.evaluator is not None:
                metrics = self.evaluator.evaluate()
                logger.info(
                    "epoch %d 验证: %s",
                    epoch,
                    " ".join(f"{k}={v:.4f}" for k, v in metrics.items()),
                )

            self.visualizer.log(
                {
                    **{f"train/epoch_{k}": v for k, v in train_metrics.items()},
                    **{f"val/{k}": v for k, v in metrics.items()},
                    **{f"train/{k}": v for k, v in self._gpu_memory().items()},
                    "epoch": epoch,
                },
                step=self.global_step,
            )

            save_checkpoint(
                self.work_dir / "last.pth",
                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),
            )
            score = metrics.get(self.monitor)
            if score is not None and score > self.best_metric:
                self.best_metric = score
                save_checkpoint(
                    self.work_dir / "best.pth",
                    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),
                )
                logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score)

        if math.isfinite(self.best_metric):
            self.visualizer.log({f"best_{self.monitor}": self.best_metric}, step=self.global_step)
        self.visualizer.finish()
        return self.best_metric