"""SimCLR 自监督预训练：在无标注的薄膜图上学一套特征，供聚类使用。

骨架仍然是项目里的纯卷积 backbone（`tiny_conv` 等），预训练完可以直接拿去
初始化检测模型。支持 cuda / mps / cpu，用 `device: auto` 自动挑最快的后端。
"""

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.metrics import AverageMeter
from engine.optim import build_optimizer, build_scheduler
from engine.visualizer import SwanLabLogger
from models.heads.projection import ProjectionHead
from utils.device import PeakMemory, resolve_amp, synchronize
from utils.logging import get_logger

logger = get_logger("engine.simclr")


def nt_xent_loss(z1: torch.Tensor, z2: torch.Tensor, temperature: float = 0.2) -> torch.Tensor:
    """NT-Xent：同一张图的两个视图互为正样本，batch 内其余 2N-2 个都是负样本。"""
    batch_size = z1.shape[0]
    z = F.normalize(torch.cat([z1, z2]), dim=1)
    similarity = z @ z.t() / temperature
    similarity.fill_diagonal_(float("-inf"))  # 排除自己和自己

    targets = torch.cat(
        [
            torch.arange(batch_size, 2 * batch_size, device=z.device),
            torch.arange(0, batch_size, device=z.device),
        ]
    )
    return F.cross_entropy(similarity, targets)


class SimCLRTrainer:
    """Args:
    backbone: 输出特征图的卷积骨架（需带 `out_channels`）。
    dataloader: 每个样本产出 2 个视图的 DataLoader（`UnlabeledImageDataset(views=2)`）。
    cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval`
        以及 `swanlab`（可视化，字段同 `train.swanlab`）。
    full_config: 记录到 SwanLab 的超参，一般直接传完整 yaml 配置。
    resume: 训练状态文件（`simclr_last.pth`）路径，从上次中断处接着练。
        与 `simclr.checkpoint` 不同：checkpoint 是"有权重就跳过训练直接聚类"，
        resume 是"载入骨架 + 投影头 + 优化器 + 调度器 + 起始 epoch 继续训练"。
    """

    def __init__(
        self,
        backbone: nn.Module,
        dataloader: DataLoader,
        cfg: Mapping[str, Any],
        device: torch.device,
        work_dir: str | Path,
        full_config: dict | None = None,
        resume: str | Path | None = None,
    ) -> None:
        self.backbone = backbone
        self.dataloader = dataloader
        self.cfg = dict(cfg)
        self.device = device
        self.work_dir = Path(work_dir)
        self.work_dir.mkdir(parents=True, exist_ok=True)

        self.epochs = int(self.cfg.get("epochs", 50))
        self.temperature = float(self.cfg.get("temperature", 0.2))
        self.log_interval = int(self.cfg.get("log_interval", 20))
        # 每隔若干 epoch 落一次盘：自监督动辄跑几十分钟，中途被打断不该把训练成果全丢掉
        self.save_interval = int(self.cfg.get("save_interval", 5))

        self.projector = ProjectionHead(
            backbone.out_channels, **dict(self.cfg.get("projection", {}))
        ).to(device)
        self.model = nn.ModuleDict({"backbone": backbone, "projector": self.projector})
        self.optimizer = build_optimizer(
            self.model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
        )
        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.global_step = 0
        self.start_epoch = 0
        if resume:
            self.load_state(resume)

        self.visualizer = SwanLabLogger(
            work_dir=self.work_dir,
            config=full_config,
            **dict(self.cfg.get("swanlab", {})),
        )

    def train_one_epoch(self, epoch: int) -> float:
        self.model.train()
        meter = AverageMeter()
        start = time.time()

        for step, batch in enumerate(self.dataloader):
            view1, view2 = (view.to(self.device, non_blocking=True) for view in batch["views"][:2])

            with torch.amp.autocast(self.device.type, enabled=self.amp):
                z1 = self.projector(self.backbone(view1))
                z2 = self.projector(self.backbone(view2))
                loss = nt_xent_loss(z1, z2, self.temperature)

            self.optimizer.zero_grad(set_to_none=True)
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()

            self.peak_memory.update()
            meter.update(float(loss.detach()), view1.shape[0])
            self.global_step += 1
            if step % self.log_interval == 0:
                self.visualizer.log(
                    {
                        "loss": float(loss.detach()),
                        "loss_avg": meter.avg,
                        "lr": self.optimizer.param_groups[0]["lr"],
                    },
                    step=self.global_step,
                    prefix="simclr/",
                )
                logger.info(
                    "simclr epoch %d [%d/%d] loss=%.4f lr=%.2e",
                    epoch,
                    step,
                    len(self.dataloader),
                    meter.avg,
                    self.optimizer.param_groups[0]["lr"],
                )

        synchronize(self.device)
        peak = self.peak_memory.peak_mb
        elapsed = time.time() - start
        self.visualizer.log(
            {
                "epoch": float(epoch),
                "epoch_loss": meter.avg,
                "epoch_seconds": elapsed,
                "gpu_mem_peak_mb": peak,
            },
            step=self.global_step,
            prefix="simclr/",
        )
        logger.info(
            "simclr epoch %d 完成 loss=%.4f 耗时 %.1fs 显存峰值 %.0f MB%s",
            epoch,
            meter.avg,
            elapsed,
            peak,
            "" if self.peak_memory.exact else "（MPS 采样下界）",
        )
        return meter.avg

    def save_backbone(self, epoch: int) -> Path:
        """只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。

        这是给下游用的产物（聚类、初始化 DefectNet）；要接着训练请用 `simclr_last.pth`。
        """
        path = self.work_dir / "simclr_backbone.pth"
        torch.save({"model": self.backbone.state_dict(), "epochs": epoch + 1}, path)
        return path

    def save_state(self, epoch: int) -> Path:
        """完整训练状态，供 `resume` 断点续训。

        比 `simclr_backbone.pth` 多了投影头、优化器、调度器与 scaler——少了这些，
        续训等于把动量和学习率进度清零，损失会先反弹一段。

        同时记下当时的优化器/调度器配置，续训时若与新配置不同会打日志说明按新的来。
        """
        path = self.work_dir / "simclr_last.pth"
        torch.save(
            {
                "backbone": self.backbone.state_dict(),
                "projector": self.projector.state_dict(),
                "optimizer": self.optimizer.state_dict(),
                "scheduler": self.scheduler.state_dict() if self.scheduler else None,
                "scaler": self.scaler.state_dict(),
                "epoch": epoch,
                "global_step": self.global_step,
                "epochs_planned": self.epochs,
                "optimizer_cfg": dict(self.cfg.get("optimizer", {})),
                "scheduler_cfg": dict(self.cfg.get("scheduler") or {}),
            },
            path,
        )
        return path

    def load_state(self, path: str | Path) -> None:
        """载入 `save_state` 写出的完整状态，从下一个 epoch 接着练。

        **配置文件是唯一真相源**：存档只提供"进度"——模型权重、优化器动量缓冲、
        已训到第几个 epoch、global_step；而学习率、weight_decay、总轮数、调度器
        一律以当前 yaml 为准。否则改了配置再续训，改动会被存档静默覆盖掉。
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"续训状态文件不存在: {path}")

        state = torch.load(path, map_location=self.device, weights_only=False)
        if "projector" not in state:
            raise ValueError(
                f"{path} 只含骨架权重，无法续训。续训要用 simclr_last.pth；"
                "若只是想跳过训练直接聚类，请改用 simclr.checkpoint。"
            )

        self.backbone.load_state_dict(state["backbone"])
        self.projector.load_state_dict(state["projector"])

        # optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay 等超参
        # 一并还原成存档时的值。先把配置构造出来的超参抄下来，加载完再按组写回去，
        # 这样只留下动量缓冲（真正的"进度"），超参仍以 yaml 为准。
        # 按组写回也保住了 build_optimizer 的分组结构（BN 与 bias 那组 weight_decay=0）。
        archived_optimizer = (state.get("optimizer_cfg") or {}).get("name")
        current_optimizer = dict(self.cfg.get("optimizer", {})).get("name")
        if archived_optimizer and current_optimizer and archived_optimizer != current_optimizer:
            # 换了优化器种类，动量缓冲的含义完全不同（AdamW 的 exp_avg 对 SGD 毫无意义），
            # 硬塞进去只会白占内存。直接不加载，从零开始积累。
            logger.warning(
                "优化器由 %s 换成了 %s，动量缓冲无法迁移、将从零积累（模型权重与训练进度不受影响）",
                archived_optimizer,
                current_optimizer,
            )
        else:
            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)

        if state.get("scaler"):
            self.scaler.load_state_dict(state["scaler"])

        self.start_epoch = int(state["epoch"]) + 1
        self.global_step = int(state.get("global_step", 0))
        self._restore_scheduler(state)
        self._log_config_changes(state)
        logger.info(
            "已载入训练状态 %s：从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs
        )
        if self.start_epoch >= self.epochs:
            logger.warning(
                "状态里已训到 epoch %d，不小于计划的 %d 轮——不会再训练。"
                "要继续请调大 simclr.epochs。",
                self.start_epoch,
                self.epochs,
            )

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

        不恢复存档里的调度器状态：它带着当时的 T_max 和 base_lrs，一旦 yaml 里改了
        总轮数或学习率就会被静默覆盖。原计划 10 轮的 cosine 跑完后 lr 已是 0，
        这时把 epochs 改成 30 接着练，沿用旧 T_max 的话 lr 仍是 0——一步都学不动。

        cosine / step / multistep 都是 epoch 的确定性函数，快进与恢复状态等价，
        所以统一走重建这一条路，没有第二种行为。
        """
        del state  # 存档里的 scheduler 状态刻意不用
        if self.scheduler is None:
            return

        # cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr，
        # 否则从 0 出发只会一路留在 0。
        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:
        """存档配置与当前配置不一致时说清楚用的是哪一份，避免"改了没生效"的困惑。"""
        changes = []
        planned = int(state.get("epochs_planned", self.epochs))
        if planned != self.epochs:
            changes.append(f"总轮数 {planned} -> {self.epochs}")
        for key, archived in (state.get("optimizer_cfg") or {}).items():
            current = dict(self.cfg.get("optimizer", {})).get(key)
            if current is not None and current != archived:
                changes.append(f"optimizer.{key} {archived} -> {current}")
        archived_scheduler = (state.get("scheduler_cfg") or {}).get("name")
        current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
        if 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 _checkpoint(self, epoch: int) -> None:
        """同时落两份：给下游用的骨架，和给续训用的完整状态。"""
        self.save_backbone(epoch)
        self.save_state(epoch)

    def fit(self) -> Path:
        """训练并保存骨架权重（不含投影头），返回骨架路径。

        每 `save_interval` 个 epoch 落一次盘（Ctrl-C 也会先落盘），
        同时写 `simclr_backbone.pth`（下游用）与 `simclr_last.pth`（续训用）。
        """
        if len(self.dataloader) < 2:
            raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")

        path = self.work_dir / "simclr_backbone.pth"
        if self.start_epoch >= self.epochs:
            self.save_backbone(self.epochs - 1)
            self.visualizer.finish()
            return path

        epoch = self.start_epoch
        try:
            for epoch in range(self.start_epoch, self.epochs):
                self.train_one_epoch(epoch)
                if self.scheduler is not None:
                    self.scheduler.step()
                if self.save_interval and (epoch + 1) % self.save_interval == 0:
                    self._checkpoint(epoch)
                    logger.info("已保存中间权重（epoch %d）: %s", epoch, path)
        except KeyboardInterrupt:
            self._checkpoint(epoch)
            logger.warning(
                "训练被中断，已保存 epoch %d 的权重: %s；续训用 --set simclr.resume=%s",
                epoch,
                path,
                self.work_dir / "simclr_last.pth",
            )
            self.visualizer.finish()
            return path

        self._checkpoint(self.epochs - 1)
        logger.info("自监督骨架已保存: %s", path)
        self.visualizer.finish()
        return path