"""SimCLR 自监督：损失的数学性质 + 聚类脚本的 simclr 分支端到端。"""

import sys

import pytest

torch = pytest.importorskip("torch")
pytest.importorskip("PIL")

import numpy as np  # noqa: E402
from PIL import Image  # noqa: E402

from engine.simclr import SimCLRTrainer, nt_xent_loss  # noqa: E402
from models.heads.projection import ProjectionHead  # noqa: E402
from models.registry import BACKBONES  # noqa: E402


def test_nt_xent_is_minimal_when_views_agree():
    """两个视图完全一致且各样本互相正交时，损失应接近理论下界 log(2B−1) 的反面——
    即正样本对独占概率质量，损失趋近 0。"""
    z = torch.eye(4)  # 4 个互相正交的样本
    aligned = nt_xent_loss(z, z.clone(), temperature=0.05)
    misaligned = nt_xent_loss(z, z.roll(1, dims=0), temperature=0.05)

    assert aligned < 0.1
    assert misaligned > aligned


def test_nt_xent_upper_bound_at_chance_level():
    """所有表征相同时无法区分正负样本，损失 = log(2B−1)。"""
    z = torch.ones(8, 16)
    loss = float(nt_xent_loss(z, z.clone(), temperature=0.2))
    assert loss == pytest.approx(float(torch.log(torch.tensor(2 * 8 - 1))), abs=1e-4)


def test_nt_xent_is_symmetric_and_temperature_sensitive():
    torch.manual_seed(0)
    z1, z2 = torch.randn(6, 8), torch.randn(6, 8)
    assert float(nt_xent_loss(z1, z2, 0.2)) == pytest.approx(float(nt_xent_loss(z2, z1, 0.2)), abs=1e-5)
    # 温度越低，随机表征的损失越大（softmax 更尖，惩罚更重）
    assert nt_xent_loss(z1, z2, 0.05) > nt_xent_loss(z1, z2, 0.5)


def test_nt_xent_gradient_pulls_positive_pair_together():
    torch.manual_seed(0)
    z1 = torch.randn(4, 8, requires_grad=True)
    z2 = torch.randn(4, 8, requires_grad=True)
    before = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())

    optimizer = torch.optim.SGD([z1, z2], lr=1.0)
    for _ in range(20):
        optimizer.zero_grad()
        nt_xent_loss(z1, z2, 0.2).backward()
        optimizer.step()

    after = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
    assert after > before


def test_projection_head_accepts_feature_map_and_vector():
    head = ProjectionHead(16, hidden_channels=32, out_channels=8)
    assert head(torch.randn(4, 16, 5, 5)).shape == (4, 8)
    assert head(torch.randn(4, 16)).shape == (4, 8)


@pytest.fixture
def tiny_image_dir(tmp_path):
    root = tmp_path / "imgs"
    root.mkdir()
    rng = np.random.default_rng(0)
    for i in range(8):
        Image.fromarray(rng.integers(0, 255, (40, 60, 3), dtype=np.uint8)).save(root / f"{i}.png")
    return root


def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
    from torch.utils.data import DataLoader

    from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate

    dataset = UnlabeledImageDataset(
        tiny_image_dir,
        transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
        views=2,
        view_transforms=[{"name": "RandomResizedCrop", "size": [32, 32], "scale": [0.5, 1.0]}],
    )
    loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
    backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})

    trainer = SimCLRTrainer(
        backbone,
        loader,
        {"epochs": 2, "amp": False, "log_interval": 100, "projection": {"hidden_channels": 16,
                                                                        "out_channels": 8}},
        torch.device("cpu"),
        tmp_path / "run",
    )
    path = trainer.fit()

    assert path.exists()
    state = torch.load(path, map_location="cpu", weights_only=False)
    # 只存骨架，不含投影头
    assert all("projector" not in key and "mlp" not in key for key in state["model"])
    backbone.load_state_dict(state["model"])  # 能被干净地加载回去


def test_simclr_logs_metrics_to_swanlab(tiny_image_dir, tmp_path, monkeypatch):
    """SwanLab 接线：按 step 记 loss/lr，按 epoch 记平均 loss、耗时与显存峰值，结束时收尾。"""
    from torch.utils.data import DataLoader

    import engine.simclr as simclr_mod
    from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate

    records: list[tuple[dict, int | None, str]] = []
    finished: list[bool] = []

    class StubLogger:
        def __init__(self, **kwargs):
            self.kwargs = kwargs

        def log(self, metrics, step=None, prefix=""):
            records.append((metrics, step, prefix))

        def finish(self):
            finished.append(True)

    monkeypatch.setattr(simclr_mod, "SwanLabLogger", StubLogger)

    dataset = UnlabeledImageDataset(
        tiny_image_dir,
        transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
        views=2,
    )
    loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
    backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})

    trainer = simclr_mod.SimCLRTrainer(
        backbone,
        loader,
        {"epochs": 1, "amp": False, "log_interval": 1, "save_interval": 0,
         "projection": {"hidden_channels": 16, "out_channels": 8},
         "swanlab": {"enabled": True, "mode": "disabled"}},
        torch.device("cpu"),
        tmp_path / "run",
        full_config={"method": "simclr"},
    )
    trainer.fit()

    assert records, "没有任何指标被记录"
    assert all(prefix == "simclr/" for _, _, prefix in records)
    step_keys = {k for metrics, _, _ in records for k in metrics}
    assert {"loss", "loss_avg", "lr"} <= step_keys
    assert {"epoch", "epoch_loss", "epoch_seconds", "gpu_mem_peak_mb"} <= step_keys
    assert all(step is not None and step > 0 for _, step, _ in records)
    assert finished == [True]  # fit() 结束时调用了 finish()

    # 完整 yaml 配置作为超参传给了 SwanLab
    assert trainer.visualizer.kwargs["config"] == {"method": "simclr"}
    assert trainer.visualizer.kwargs["mode"] == "disabled"


def test_simclr_saves_intermediate_checkpoints(tiny_image_dir, tmp_path):
    """中途落盘：跑到一半被打断也要能拿到权重（此前只在全部 epoch 跑完后才存）。"""
    from torch.utils.data import DataLoader

    from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate

    dataset = UnlabeledImageDataset(
        tiny_image_dir,
        transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
        views=2,
    )
    loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
    backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
    work_dir = tmp_path / "run"

    trainer = SimCLRTrainer(
        backbone,
        loader,
        {"epochs": 4, "amp": False, "log_interval": 100, "save_interval": 2,
         "projection": {"hidden_channels": 16, "out_channels": 8}},
        torch.device("cpu"),
        work_dir,
    )
    # 手动走前两个 epoch，模拟"训练中途"
    trainer.train_one_epoch(0)
    trainer.train_one_epoch(1)
    path = trainer.save_backbone(1)

    assert path.exists()
    assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 2

    trainer.fit()  # 完整跑完后 epochs 应更新
    assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 4


def _make_trainer(tiny_image_dir, work_dir, epochs=4, **cfg_extra):
    from torch.utils.data import DataLoader

    from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate

    dataset = UnlabeledImageDataset(
        tiny_image_dir,
        transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
        views=2,
    )
    loader = DataLoader(dataset, batch_size=4, shuffle=False, drop_last=True, collate_fn=unlabeled_collate)
    torch.manual_seed(0)
    backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
    cfg = {
        "epochs": epochs,
        "amp": False,
        "log_interval": 100,
        "save_interval": 1,
        "projection": {"hidden_channels": 16, "out_channels": 8},
        "scheduler": {"name": "cosine"},
        "swanlab": {"enabled": False},
        **cfg_extra,
    }
    return SimCLRTrainer(backbone, loader, cfg, torch.device("cpu"), work_dir,
                         resume=cfg_extra.pop("resume", None))


def test_save_state_contains_full_training_state(tiny_image_dir, tmp_path):
    trainer = _make_trainer(tiny_image_dir, tmp_path / "run", epochs=1)
    trainer.fit()

    state_path = tmp_path / "run" / "simclr_last.pth"
    assert state_path.exists()
    state = torch.load(state_path, map_location="cpu", weights_only=False)
    # 续训需要的东西一个都不能少，否则动量和学习率进度会被清零
    assert {"backbone", "projector", "optimizer", "scheduler", "scaler", "epoch", "global_step"} <= set(state)
    assert state["epoch"] == 0


def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path):
    """resume 是接着练，不是从头再来。"""
    work_dir = tmp_path / "run"
    first = _make_trainer(tiny_image_dir, work_dir, epochs=2)
    first.fit()
    steps_after_2 = first.global_step
    lr_after_2 = first.optimizer.param_groups[0]["lr"]

    resumed = SimCLRTrainer(
        BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
        first.dataloader,
        {"epochs": 5, "amp": False, "log_interval": 100, "save_interval": 1,
         "projection": {"hidden_channels": 16, "out_channels": 8},
         "scheduler": {"name": "cosine"}, "swanlab": {"enabled": False}},
        torch.device("cpu"),
        work_dir,
        resume=work_dir / "simclr_last.pth",
    )
    assert resumed.start_epoch == 2                    # 从第 3 轮开始
    assert resumed.global_step == steps_after_2        # step 计数接上
    assert resumed.optimizer.param_groups[0]["lr"] == pytest.approx(lr_after_2)  # 学习率进度接上

    resumed.fit()
    state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False)
    assert state["epoch"] == 4                         # 一路练到第 5 轮
    assert resumed.global_step > steps_after_2


def test_resume_rejects_backbone_only_file(tiny_image_dir, tmp_path):
    """simclr_backbone.pth 只有骨架，不能当续训状态用，要给出可操作的报错。"""
    work_dir = tmp_path / "run"
    _make_trainer(tiny_image_dir, work_dir, epochs=1).fit()

    with pytest.raises(ValueError, match="只含骨架权重"):
        SimCLRTrainer(
            BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
            _make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
            {"epochs": 2, "amp": False, "projection": {"hidden_channels": 16, "out_channels": 8},
             "swanlab": {"enabled": False}},
            torch.device("cpu"),
            work_dir,
            resume=work_dir / "simclr_backbone.pth",
        )

    with pytest.raises(FileNotFoundError, match="不存在"):
        _make_trainer(tiny_image_dir, work_dir, epochs=1, resume=work_dir / "nope.pth")


def test_resume_beyond_planned_epochs_does_not_train(tiny_image_dir, tmp_path):
    work_dir = tmp_path / "run"
    _make_trainer(tiny_image_dir, work_dir, epochs=3).fit()

    resumed = SimCLRTrainer(
        BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
        _make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
        {"epochs": 3, "amp": False, "projection": {"hidden_channels": 16, "out_channels": 8},
         "swanlab": {"enabled": False}},
        torch.device("cpu"),
        work_dir,
        resume=work_dir / "simclr_last.pth",
    )
    steps_before = resumed.global_step
    resumed.fit()
    assert resumed.global_step == steps_before  # 一步都没再训


def test_resume_and_checkpoint_are_mutually_exclusive(tiny_image_dir, tmp_path, monkeypatch):
    sys.path.insert(0, "scripts")
    import cluster_unlabeled as cu

    with pytest.raises(ValueError, match="只能二选一"):
        cu.build_encoder(
            {"method": "simclr",
             "simclr": {"resume": "a.pth", "checkpoint": "b.pth"},
             "dataset": {"name": "unlabeled", "root": str(tiny_image_dir)}},
            torch.device("cpu"),
            tmp_path,
        )


def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monkeypatch):
    """method=simclr 全流程：训练 -> 存权重 -> 抽特征 -> 聚类 -> 产出。"""
    sys.path.insert(0, "scripts")
    import cluster_unlabeled as cu

    work_dir = tmp_path / "run"
    monkeypatch.setattr(
        sys, "argv",
        ["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
         "--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
         "device=cpu", "method=simclr", "loader.num_workers=0", "loader.batch_size=4",
         "simclr.epochs=1", "simclr.batch_size=4", "simclr.amp=False",
         "simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
         "simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
         "num_clusters=2", "n_init=2"],
    )

    assert cu.main() == 0
    assert (work_dir / "simclr_backbone.pth").exists()
    assert (work_dir / "assignments.csv").exists()
    assert (work_dir / "summary.json").exists()
    assert (work_dir / "clusters" / "cluster_00").is_dir()


def test_simclr_runs_on_the_resolved_device(tiny_image_dir, tmp_path, monkeypatch):
    """method=simclr 在本机自动选中的后端上（MPS / CUDA / CPU）都应能跑通。"""
    sys.path.insert(0, "scripts")
    import cluster_unlabeled as cu

    from utils.device import describe_device, resolve_device

    device = resolve_device("auto")
    work_dir = tmp_path / "auto_run"
    monkeypatch.setattr(
        sys, "argv",
        ["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
         "--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
         "method=simclr", "loader.num_workers=0", "loader.batch_size=4",
         "simclr.epochs=1", "simclr.batch_size=4",
         "simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
         "simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
         "num_clusters=2", "n_init=2"],
    )

    assert cu.main() == 0, f"在 {describe_device(device)} 上跑失败"
    assert (work_dir / "simclr_backbone.pth").exists()