"""卷积骨架与分类模型。"""

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.heads.projection import ProjectionHead  # noqa: E402


@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"})


def test_tiny_conv_shapes():
    from models.registry import BACKBONES

    backbone = BACKBONES.build(
        {"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]}
    ).eval()
    with torch.no_grad():
        feat = backbone(torch.randn(2, 3, 64, 64))
    assert feat.shape[:2] == (2, 32)
    assert backbone.out_channels == 32
    # stem 下采样 2 倍，其后每个 stage(除第一个)各下采样 2 倍
    assert backbone.out_stride == 8


@pytest.mark.parametrize("arch", ["resnet18", "shufflenet_v2_x1_0", "mobilenet_v3_small"])
def test_alternative_backbones_drive_simclr(arch):
    """换骨架后 SimCLR 的投影头要能接上（out_channels 必须准）。"""
    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_classifier_shapes_and_convolution_only():
    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())