import pytest
import torch

from giant.model.layers import build_mlp_head


def test_build_mlp_head_depth_1_is_bare_linear():
    head = build_mlp_head(8, 4, hidden=16, depth=1)
    assert len(head) == 1
    assert isinstance(head[0], torch.nn.Linear)
    assert head[0].in_features == 8
    assert head[0].out_features == 4
    out = head(torch.randn(3, 8))
    assert out.shape == (3, 4)


def test_build_mlp_head_depth_2_matches_pre_gitea_36_shape():
    head = build_mlp_head(8, 4, hidden=16, depth=2)
    assert len(head) == 3
    assert isinstance(head[0], torch.nn.Linear)
    assert head[0].in_features == 8
    assert head[0].out_features == 16
    assert isinstance(head[1], torch.nn.SiLU)
    assert isinstance(head[2], torch.nn.Linear)
    assert head[2].in_features == 16
    assert head[2].out_features == 4
    out = head(torch.randn(5, 8))
    assert out.shape == (5, 4)


def test_build_mlp_head_depth_3_has_extra_hidden_layer():
    head = build_mlp_head(8, 4, hidden=16, depth=3)
    assert len(head) == 5
    widths = [(m.in_features, m.out_features) for m in head if isinstance(m, torch.nn.Linear)]
    assert widths == [(8, 16), (16, 16), (16, 4)]
    out = head(torch.randn(2, 8))
    assert out.shape == (2, 4)


def test_build_mlp_head_depth_0_raises():
    with pytest.raises(ValueError, match="depth"):
        build_mlp_head(8, 4, hidden=16, depth=0)