import pytest
import torch

from mini_transformer.modules.lm_head import LMHead

# -----------------------
# Constructor checks
# -----------------------


def test_ctor_type_and_value_checks():
    with pytest.raises(TypeError):
        LMHead("64", 10)
    with pytest.raises(TypeError):
        LMHead(64, "10")
    with pytest.raises(ValueError):
        LMHead(0, 10)
    with pytest.raises(ValueError):
        LMHead(64, 0)


def test_ctor_happy_path():
    lm = LMHead(32, 1000)
    assert lm.d_model == 32
    assert lm.vocab_size == 1000
    assert lm.fc.weight.shape == (1000, 32)


# -----------------------
# Forward checks (3D only)
# -----------------------


def test_forward_type_and_shape_errors():
    lm = LMHead(16, 50)
    with pytest.raises(TypeError):
        lm("not a tensor")
    with pytest.raises(ValueError):
        lm(torch.randn(2, 3, 4, 5))  # rank 4 not supported
    with pytest.raises(ValueError):
        lm(torch.randn(5, 16))  # rank 2 not supported
    with pytest.raises(ValueError):
        lm(torch.randn(2, 8, 15))  # last dim != d_model


def test_forward_happy_path_and_zero_len():
    device = (
        torch.device(f"cuda:{torch.cuda.current_device()}")
        if torch.cuda.is_available()
        else torch.device("cpu")
    )
    lm = LMHead(24, 101).to(device)

    # Non-empty
    x = torch.randn(3, 7, 24, device=device, dtype=torch.float32)
    y = lm(x)
    assert y.shape == (3, 7, 101)
    assert y.device == device and y.dtype == torch.float32

    # Zero-length sequence allowed
    x0 = torch.randn(2, 0, 24, device=device, dtype=torch.float32)
    y0 = lm(x0)
    assert y0.shape == (2, 0, 101)


def test_gradients_flow():
    lm = LMHead(8, 40)
    x = torch.randn(3, 6, 8, requires_grad=True)
    y = lm(x)
    loss = y.pow(2).mean()
    loss.backward()
    assert x.grad is not None
    assert torch.isfinite(x.grad).all()