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

骨架仍然是项目里的纯卷积 backbone（`tiny_conv` 等），预训练完可以直接拿去
初始化检测模型。属于"训练"，因此调用方必须先过 `utils.device.require_cuda()`。
"""

from __future__ import annotations

import time
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 models.heads.projection import ProjectionHead
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` 等。
    """

    def __init__(
        self,
        backbone: nn.Module,
        dataloader: DataLoader,
        cfg: Mapping[str, Any],
        device: torch.device,
        work_dir: str | Path,
    ) -> 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))

        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 = bool(self.cfg.get("amp", True)) and device.type == "cuda"
        self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)

    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("cuda", 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()

            meter.update(float(loss.detach()), view1.shape[0])
            if step % self.log_interval == 0:
                logger.info(
                    "simclr epoch %d [%d/%d] loss=%.4f lr=%.2e",
                    epoch,
                    step,
                    len(self.dataloader),
                    meter.avg,
                    self.optimizer.param_groups[0]["lr"],
                )

        logger.info("simclr epoch %d 完成 loss=%.4f 耗时 %.1fs", epoch, meter.avg, time.time() - start)
        return meter.avg

    def fit(self) -> Path:
        """训练并保存骨架权重（不含投影头），返回权重路径。"""
        if len(self.dataloader) < 2:
            raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")

        for epoch in range(self.epochs):
            self.train_one_epoch(epoch)
            if self.scheduler is not None:
                self.scheduler.step()

        path = self.work_dir / "simclr_backbone.pth"
        torch.save({"model": self.backbone.state_dict(), "epochs": self.epochs}, path)
        logger.info("自监督骨架已保存: %s", path)
        return path