import pytest

torch = pytest.importorskip("torch")

from models import build_model  # 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