import pytest

torch = pytest.importorskip("torch")

from models import build_model  # noqa: E402
from models.backbones.resnet import _SUPPORTED as _SUPPORTED_ARCHS  # noqa: E402
from models.postprocess import decode_heatmap  # noqa: E402


@pytest.fixture
def model_cfg():
    return {
        "name": "defect_net",
        "num_classes": 3,
        "backbone": {"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]},
        "neck": {"name": "upsample", "channels": 16, "num_upsamples": 2},
        "head": {"hidden_channels": 16},
    }


def test_forward_shapes_and_stride(model_cfg):
    model = build_model(model_cfg).eval()
    # stem(2) * 2 个下采样 stage = 8，再上采样 4 倍 => stride 2
    assert model.out_stride == 2

    with torch.no_grad():
        out = model(torch.randn(2, 3, 64, 64))

    size = 64 // model.out_stride
    assert out["heatmap"].shape == (2, 3, size, size)
    assert out["wh"].shape == (2, 2, size, size)
    assert out["offset"].shape == (2, 2, size, size)
    assert out["logits"].shape == (2, 3)


def test_backbone_is_convolution_only(model_cfg):
    model = build_model(model_cfg)
    forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM, torch.nn.GRU)
    assert not any(isinstance(m, forbidden) for m in model.modules())


def test_without_classification_head(model_cfg):
    model = build_model({**model_cfg, "with_classification": False}).eval()
    with torch.no_grad():
        out = model(torch.randn(1, 3, 64, 64))
    assert "logits" not in out


def test_decode_heatmap_recovers_planted_box():
    heatmap = torch.full((1, 2, 32, 32), -10.0)
    heatmap[0, 1, 8, 10] = 10.0
    wh = torch.zeros(1, 2, 32, 32)
    wh[0, :, 8, 10] = torch.tensor([4.0, 6.0])
    offset = torch.zeros(1, 2, 32, 32)

    (result,) = decode_heatmap(heatmap, wh, offset, stride=4, topk=5, score_threshold=0.5)

    assert result["labels"].tolist() == [1]
    assert result["boxes"][0].tolist() == pytest.approx([32.0, 20.0, 48.0, 44.0])


def test_losses_decrease_on_a_single_batch(model_cfg):
    from engine.losses import DetectionLoss

    torch.manual_seed(0)
    model = build_model(model_cfg)
    criterion = DetectionLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

    images = torch.randn(2, 3, 64, 64)
    size = 64 // model.out_stride
    target = {
        "heatmap": torch.zeros(2, 3, size, size),
        "wh": torch.zeros(2, 2, size, size),
        "offset": torch.zeros(2, 2, size, size),
        "mask": torch.zeros(2, size, size),
    }
    target["heatmap"][0, 1, 5, 5] = 1.0
    target["mask"][0, 5, 5] = 1.0
    target["wh"][0, :, 5, 5] = 3.0
    batch = {"target": target, "image_labels": torch.zeros(2, 3)}

    first = last = None
    for step in range(5):
        losses = criterion(model(images), batch)
        optimizer.zero_grad()
        losses["total"].backward()
        optimizer.step()
        total = float(losses["total"].detach())
        first = total if step == 0 else first
        last = total

    assert last < first


@pytest.mark.parametrize("arch", sorted(_SUPPORTED_ARCHS))
def test_every_supported_torchvision_arch_actually_works(arch):
    """清单里列出的每个骨架都必须真能前向，且声明的 out_channels 与实际一致。

    回归测试：shufflenet_v2_x1_0 曾被列为支持，但它没有 `.features` 属性，
    构造时直接 AttributeError。
    """
    from models.registry import BACKBONES

    backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval()
    with torch.no_grad():
        feat = backbone(torch.randn(1, 3, 64, 64))

    assert feat.dim() == 4
    assert feat.shape[1] == backbone.out_channels == _SUPPORTED_ARCHS[arch]
    assert backbone.out_stride == 32


def test_unsupported_arch_is_rejected():
    from models.registry import BACKBONES

    with pytest.raises(ValueError, match="不支持的骨架"):
        BACKBONES.build({"name": "torchvision", "arch": "vit_b_16"})


@pytest.mark.parametrize("arch", ["resnet18", "shufflenet_v2_x1_0", "mobilenet_v3_small"])
def test_alternative_backbones_drive_simclr(arch):
    """换骨架后 SimCLR 的投影头要能接上（out_channels 必须准）。"""
    from models.heads.projection import ProjectionHead
    from models.registry import BACKBONES

    backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval()
    head = ProjectionHead(backbone.out_channels, 32, 16)
    with torch.no_grad():
        assert head(backbone(torch.randn(2, 3, 64, 64))).shape == (2, 16)


def test_alternative_backbone_drives_defect_net():
    """检测模型换骨架同样要能跑通（neck 的 in_channels 由 backbone 决定）。"""
    model = build_model(
        {
            "name": "defect_net",
            "num_classes": 3,
            "backbone": {"name": "torchvision", "arch": "shufflenet_v2_x1_0", "pretrained": False},
            "neck": {"name": "upsample", "channels": 16, "num_upsamples": 3},
            "head": {"hidden_channels": 16},
        }
    ).eval()
    assert model.out_stride == 4
    with torch.no_grad():
        out = model(torch.randn(1, 3, 64, 64))
    assert out["heatmap"].shape == (1, 3, 16, 16)