"""分类数据集、模型与训练引擎。"""

import sys

import pytest

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

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

from datasets import FolderClassificationDataset, classification_collate  # noqa: E402
from engine.classification import (  # noqa: E402
    ClassificationMetric,
    ClassificationTrainer,
    class_weights,
)
from models import build_model  # noqa: E402

TRANSFORMS = [{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}]


@pytest.fixture
def dataset_root(tmp_path):
    """三个类别，样本数刻意不均衡（20 / 6 / 2）。"""
    rng = np.random.default_rng(0)
    root = tmp_path / "dataset"          # 单独一层：输出目录不能落在数据集根下，
    root.mkdir()                         # 否则会被当成一个新类别扫进来
    for name, count in (("normal", 20), ("grease", 6), ("others", 2)):
        directory = root / name
        directory.mkdir()
        for i in range(count):
            Image.fromarray(rng.integers(0, 255, (18, 24, 3), dtype=np.uint8)).save(
                directory / f"{i}.png"
            )
    return root


def test_classes_come_from_subdirectory_names(dataset_root):
    dataset = FolderClassificationDataset(dataset_root, split="all", transforms=TRANSFORMS)
    assert dataset.classes == ["grease", "normal", "others"]  # 按目录名排序
    assert len(dataset) == 28
    assert dataset.class_counts() == [6, 20, 2]


def test_split_is_stratified_and_deterministic(dataset_root):
    train = FolderClassificationDataset(dataset_root, split="train", transforms=TRANSFORMS)
    val = FolderClassificationDataset(dataset_root, split="val", transforms=TRANSFORMS)

    assert len(train) + len(val) == 28
    # 逐类切分：即使只有 2 张的类别也要在验证集里留下样本
    assert all(count > 0 for count in val.class_counts()), val.class_counts()
    assert all(count > 0 for count in train.class_counts())

    again = FolderClassificationDataset(dataset_root, split="val", transforms=TRANSFORMS)
    assert [p for p, _ in val.samples] == [p for p, _ in again.samples]

    # 训练与验证不能有交集
    assert not {p for p, _ in train.samples} & {p for p, _ in val.samples}


def test_different_seed_changes_the_split(dataset_root):
    a = FolderClassificationDataset(dataset_root, split="val", seed=1, transforms=TRANSFORMS)
    b = FolderClassificationDataset(dataset_root, split="val", seed=2, transforms=TRANSFORMS)
    assert [p for p, _ in a.samples] != [p for p, _ in b.samples]


def test_exclude_and_explicit_classes(dataset_root):
    dataset = FolderClassificationDataset(
        dataset_root, split="all", transforms=TRANSFORMS, exclude=["others"]
    )
    assert dataset.classes == ["grease", "normal"]

    ordered = FolderClassificationDataset(
        dataset_root, split="all", transforms=TRANSFORMS, classes=["others", "normal"]
    )
    assert ordered.class_to_idx == {"others": 0, "normal": 1}

    with pytest.raises(FileNotFoundError, match="不存在"):
        FolderClassificationDataset(dataset_root, classes=["nope"], transforms=TRANSFORMS)


def test_sample_and_collate_shapes(dataset_root):
    dataset = FolderClassificationDataset(dataset_root, split="all", transforms=TRANSFORMS)
    sample = dataset[0]
    assert sample["image"].shape == (3, 32, 32)
    assert isinstance(sample["label"], int)
    assert sample["meta"]["class"] == dataset.classes[sample["label"]]

    batch = classification_collate([dataset[0], dataset[1]])
    assert batch["images"].shape == (2, 3, 32, 32)
    assert batch["labels"].shape == (2,) and batch["labels"].dtype == torch.long


def test_class_weights_favour_rare_classes():
    weights = class_weights([100, 10, 1])
    assert weights[2] > weights[1] > weights[0]
    assert float(weights.mean()) == pytest.approx(1.0)
    assert class_weights([1, 2], mode="none") is None
    with pytest.raises(ValueError, match="未知 class_weight"):
        class_weights([1, 2], mode="sqrt")


def test_metric_matches_hand_computed_values():
    metric = ClassificationMetric(["a", "b"])
    # 真实 a,a,b,b -> 预测 a,b,b,b：a 召回 0.5，b 精确 2/3
    metric.update(torch.tensor([[9.0, 0.0], [0.0, 9.0], [0.0, 9.0], [0.0, 9.0]]),
                  torch.tensor([0, 0, 1, 1]))
    report = metric.per_class()
    assert report["a"] == pytest.approx({"precision": 1.0, "recall": 0.5,
                                         "f1": 2 / 3, "support": 2})
    assert report["b"]["precision"] == pytest.approx(2 / 3)

    metrics = metric.compute()
    assert metrics["accuracy"] == pytest.approx(0.75)
    assert metrics["macro_f1"] == pytest.approx((2 / 3 + 0.8) / 2)


def test_macro_f1_punishes_majority_only_prediction():
    """全预测成多数类：准确率还行，但 macro-F1 很低——这正是选它当主指标的原因。"""
    metric = ClassificationMetric(["major", "minor"])
    logits = torch.tensor([[9.0, 0.0]]).repeat(100, 1)
    targets = torch.tensor([0] * 90 + [1] * 10)
    metric.update(logits, targets)

    metrics = metric.compute()
    assert metrics["accuracy"] == pytest.approx(0.9)
    assert metrics["macro_f1"] < 0.5
    assert metrics["balanced_accuracy"] == pytest.approx(0.5)


def test_classifier_is_convolution_only_and_shapes_match():
    model = build_model({
        "name": "classifier", "num_classes": 4,
        "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
        "classification": {"dropout": 0.1},
    }).eval()
    with torch.no_grad():
        assert model(torch.randn(3, 3, 64, 64)).shape == (3, 4)
    forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM)
    assert not any(isinstance(m, forbidden) for m in model.modules())


def test_init_from_loads_self_supervised_backbone(tmp_path):
    """自监督骨架能直接拿来初始化分类模型——这是聚类阶段产物的下游用途。"""
    from models.registry import BACKBONES

    backbone_cfg = {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}
    pretrained = BACKBONES.build(dict(backbone_cfg))
    with torch.no_grad():
        for param in pretrained.parameters():
            param.fill_(0.5)
    path = tmp_path / "simclr_backbone.pth"
    torch.save({"model": pretrained.state_dict(), "epochs": 40}, path)

    model = build_model({"name": "classifier", "num_classes": 3,
                         "backbone": dict(backbone_cfg), "init_from": str(path)})
    for param in model.backbone.parameters():
        assert torch.allclose(param, torch.full_like(param, 0.5))

    with pytest.raises(FileNotFoundError):
        build_model({"name": "classifier", "num_classes": 3,
                     "backbone": dict(backbone_cfg), "init_from": str(tmp_path / "nope.pth")})


def test_frozen_backbone_only_trains_the_head():
    model = build_model({
        "name": "classifier", "num_classes": 3,
        "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
        "frozen_backbone": True,
    })
    assert not any(p.requires_grad for p in model.backbone.parameters())
    assert all(p.requires_grad for p in model.head.parameters())


def _trainer(dataset_root, tmp_path, **cfg_extra):
    from torch.utils.data import DataLoader

    datasets = {
        split: FolderClassificationDataset(dataset_root, split=split, transforms=TRANSFORMS)
        for split in ("train", "val")
    }
    loaders = {
        split: DataLoader(ds, batch_size=4, shuffle=(split == "train"),
                          collate_fn=classification_collate)
        for split, ds in datasets.items()
    }
    model = build_model({"name": "classifier", "num_classes": len(datasets["train"].classes),
                         "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}})
    cfg = {"epochs": 2, "amp": False, "log_interval": 999, "scheduler": {"name": "cosine"},
           "optimizer": {"name": "adamw", "lr": 1e-3}, "swanlab": {"enabled": False}, **cfg_extra}
    return ClassificationTrainer(
        model, loaders["train"], loaders["val"], cfg, torch.device("cpu"), tmp_path / "run",
        classes=datasets["train"].classes, train_counts=datasets["train"].class_counts(),
    )


def test_training_runs_and_saves_best_by_macro_f1(dataset_root, tmp_path):
    trainer = _trainer(dataset_root, tmp_path)
    best = trainer.fit()

    assert (tmp_path / "run" / "best.pth").exists()
    assert (tmp_path / "run" / "last.pth").exists()
    assert best > float("-inf")
    assert trainer.monitor == "macro_f1"   # 不均衡数据下的默认主指标

    metrics = trainer.evaluate()
    assert {"accuracy", "macro_f1", "weighted_f1", "balanced_accuracy", "loss"} <= set(metrics)


def test_resume_takes_hyperparams_from_config(dataset_root, tmp_path):
    """与检测/自监督一致：检查点只提供进度，超参以配置为准。"""
    trainer = _trainer(dataset_root, tmp_path, epochs=2,
                       optimizer={"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
    trainer.fit()

    resumed = _trainer(dataset_root, tmp_path, epochs=6,
                       optimizer={"name": "adamw", "lr": 5e-4, "weight_decay": 0.05})
    resumed.resume(tmp_path / "run" / "last.pth")

    assert resumed.start_epoch == 2
    groups = resumed.optimizer.param_groups
    assert groups[0]["initial_lr"] == pytest.approx(5e-4)
    assert groups[0]["weight_decay"] == pytest.approx(0.05)
    assert resumed.scheduler.T_max == 6


def test_cli_train_and_eval_only(dataset_root, tmp_path, monkeypatch):
    import classify

    config = tmp_path / "cfg.yaml"
    config.write_text(
        f"""
work_dir: {tmp_path / "run_cli"}
seed: 42
device: cpu
data:
  root: {dataset_root}
  val_ratio: 0.25
  exclude: []
  classes: null
  train_transforms: [{{name: Resize, size: [32, 32]}}, {{name: ToTensor}}]
  val_transforms: [{{name: Resize, size: [32, 32]}}, {{name: ToTensor}}]
  loader: {{batch_size: 4, num_workers: 0}}
model:
  name: classifier
  backbone: {{name: tiny_conv, widths: [8, 16], depths: [1, 1]}}
  classification: {{dropout: 0.0}}
  init_from: null
  frozen_backbone: false
train:
  epochs: 1
  amp: false
  log_interval: 999
  monitor: macro_f1
  class_weight: balanced
  optimizer: {{name: adamw, lr: 0.001}}
  scheduler: {{name: cosine}}
  swanlab: {{enabled: false}}
""",
        encoding="utf-8",
    )

    monkeypatch.setattr(sys, "argv", ["classify.py", "--config", str(config)])
    assert classify.main() == 0
    assert (tmp_path / "run_cli" / "eval_report.json").exists()

    monkeypatch.setattr(
        sys, "argv",
        ["classify.py", "--config", str(config), "--eval-only",
         "--checkpoint", str(tmp_path / "run_cli" / "best.pth")],
    )
    assert classify.main() == 0