"""多类(multilabel) / 单类(multiclass) 由 yaml 的 label_mode 决定，两条路径都要能跑通。"""

import pytest

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

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

from datasets import FilmDefectDataset, collate_fn  # noqa: E402
from engine.losses import DetectionLoss  # noqa: E402
from engine.metrics import SingleLabelMetric  # noqa: E402
from models import build_model  # noqa: E402
from models.defect_net import classification_outputs  # noqa: E402

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


def make_model(label_mode: str, num_classes: int = 2):
    return build_model(
        {
            "name": "defect_net",
            "num_classes": num_classes,
            "label_mode": label_mode,
            "backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
            "neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
            "head": {"hidden_channels": 8},
        }
    )


@pytest.fixture
def dataset_factory(tmp_path):
    """YOLO 目录：a.png 上 scratch 1 个框、bubble 2 个框；b.png 无缺陷。"""
    images = tmp_path / "images"
    labels = tmp_path / "labels"
    images.mkdir()
    labels.mkdir()
    for name in ("a.png", "b.png"):
        Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(images / name)

    (labels / "a.txt").write_text(
        "0 0.1875 0.1875 0.25 0.25\n"  # scratch
        "1 0.5000 0.1875 0.25 0.25\n"  # bubble
        "1 0.2500 0.5000 0.25 0.25\n",  # bubble
        encoding="utf-8",
    )
    (labels / "b.txt").write_text("", encoding="utf-8")

    def build(label_mode: str):
        return FilmDefectDataset(
            img_dir=images,
            classes=CLASSES,
            transforms=TRANSFORMS,
            stride=2,
            label_mode=label_mode,
        )

    return build


def test_multilabel_dataset_yields_multi_hot(dataset_factory):
    dataset = dataset_factory("multilabel")
    assert dataset[0]["image_label"].tolist() == [1.0, 1.0]
    assert dataset[1]["image_label"].tolist() == [0.0, 0.0]


def test_multiclass_dataset_yields_exclusive_index(dataset_factory):
    dataset = dataset_factory("multiclass")
    # bubble 出现 2 次、scratch 1 次 -> 取 bubble(下标 1)
    assert dataset[0]["image_label"].item() == 1
    assert dataset[0]["image_label"].dtype == torch.long
    # 无缺陷 -> 背景下标 = len(classes)
    assert dataset[1]["image_label"].item() == len(CLASSES)


def test_collate_shapes_differ_by_mode(dataset_factory):
    multilabel = collate_fn([dataset_factory("multilabel")[i] for i in range(2)])
    multiclass = collate_fn([dataset_factory("multiclass")[i] for i in range(2)])
    assert multilabel["image_labels"].shape == (2, len(CLASSES))
    assert multiclass["image_labels"].shape == (2,)
    # 定位监督与 label_mode 无关
    assert torch.equal(multilabel["target"]["heatmap"], multiclass["target"]["heatmap"])


def test_classification_head_width():
    assert classification_outputs(5, "multilabel") == 5
    assert classification_outputs(5, "multiclass") == 6
    assert make_model("multilabel").cls_head.num_outputs == 2
    assert make_model("multiclass").cls_head.num_outputs == 3


def test_invalid_label_mode_raises():
    with pytest.raises(ValueError, match="label_mode"):
        make_model("single")
    with pytest.raises(ValueError, match="label_mode"):
        DetectionLoss(label_mode="single")


@pytest.mark.parametrize("label_mode", ["multilabel", "multiclass"])
def test_end_to_end_loss_backward(dataset_factory, label_mode):
    model = make_model(label_mode)
    criterion = DetectionLoss(label_mode=label_mode)
    dataset = dataset_factory(label_mode)
    batch = collate_fn([dataset[0], dataset[1]])

    losses = criterion(model(batch["images"]), batch)
    losses["total"].backward()

    assert set(losses) == {"heatmap", "wh", "offset", "cls", "total"}
    assert torch.isfinite(losses["total"])
    assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.cls_head.parameters())


def test_single_label_metric_counts_only_present_classes():
    metric = SingleLabelMetric(num_outputs=3)
    logits = torch.tensor([[9.0, 0.0, 0.0], [0.0, 9.0, 0.0], [9.0, 0.0, 0.0]])
    metric.update(logits, torch.tensor([0, 1, 1]))
    result = metric.compute()

    assert result["cls_accuracy"] == pytest.approx(2 / 3)
    # 类别 2 没出现，不参与 macro 平均
    assert result["cls_f1"] == pytest.approx((2 / 3 + 2 / 3) / 2)