import pytest
import torch

from mini_transformer.modules.attention import MultiHeadAttention


# -----------------------
# Helpers
# -----------------------
def _rand(B=2, S=5, D=12, device="cpu", dtype=torch.float32):
    return torch.randn(B, S, D, device=device, dtype=dtype)


def _pad_mask(batch: int, heads: int, seq_len: int, device) -> torch.Tensor:
    return torch.zeros(batch, heads, 1, seq_len, dtype=torch.bool, device=device)


# =======================
# Constructor checks
# =======================


def test_ctor_type_checks():
    with pytest.raises(TypeError):
        MultiHeadAttention(d_model="64", num_heads=4, dropout_rate=0.1)
    with pytest.raises(TypeError):
        MultiHeadAttention(d_model=64, num_heads=4.0, dropout_rate=0.1)
    with pytest.raises(TypeError):
        MultiHeadAttention(d_model=64, num_heads=4, dropout_rate="0.1")


def test_ctor_value_checks():
    with pytest.raises(ValueError):
        MultiHeadAttention(d_model=0, num_heads=4, dropout_rate=0.1)
    with pytest.raises(ValueError):
        MultiHeadAttention(d_model=64, num_heads=0, dropout_rate=0.1)
    with pytest.raises(ValueError):
        MultiHeadAttention(d_model=63, num_heads=4, dropout_rate=0.1)  # not divisible


def test_ctor_happy_path():
    mha = MultiHeadAttention(64, 8, 0.1)
    assert mha.d_model == 64 and mha.num_heads == 8 and mha.d_head == 8


# =======================
# Forward: type/shape checks
# =======================


def test_forward_requires_tensors_and_mask_type():
    mha = MultiHeadAttention(32, 4, 0.1)
    q, k, v = _rand(2, 5, 32), _rand(2, 5, 32), _rand(2, 5, 32)
    mask = _pad_mask(2, mha.num_heads, 5, q.device)

    with pytest.raises(TypeError):
        mha("q", k, v, mask, mask)
    with pytest.raises(TypeError):
        mha(q, "k", v, mask, mask)
    with pytest.raises(TypeError):
        mha(q, k, "v", mask, mask)
    with pytest.raises(TypeError):
        mha(q, k, v, "bad", mask)
    with pytest.raises(TypeError):
        mha(q, k, v, mask, "bad")


def test_forward_rank_and_lastdim_checks():
    mha = MultiHeadAttention(32, 4, 0.1)
    q = torch.randn(2, 5, 32)
    k = torch.randn(2, 5, 32)
    v = torch.randn(2, 5, 32)
    mask = _pad_mask(2, mha.num_heads, 5, q.device)

    with pytest.raises(ValueError):
        bad_q = q.unsqueeze(0)
        bad_q_mask = _pad_mask(bad_q.shape[0], mha.num_heads, bad_q.shape[1], bad_q.device)
        mha(bad_q, k, v, bad_q_mask, mask)
    with pytest.raises(ValueError):
        mha(q, k.view(10, 32), v, mask, mask)
    with pytest.raises(ValueError):
        mha(q[..., :16], k, v, mask, mask)
    with pytest.raises(ValueError):
        mha(q, k[..., :16], v, mask, mask)
    with pytest.raises(ValueError):
        mha(q, k, v[..., :16], mask, mask)


def test_forward_batch_and_seq_mismatch_checks():
    mha = MultiHeadAttention(24, 3, 0.1)
    q = torch.randn(2, 5, 24)
    k = torch.randn(3, 5, 24)  # batch mismatch
    v = torch.randn(2, 5, 24)  # seq mismatch with k
    q_mask = _pad_mask(q.shape[0], mha.num_heads, q.shape[1], q.device)
    k_mask = _pad_mask(k.shape[0], mha.num_heads, k.shape[1], k.device)

    with pytest.raises(ValueError):
        mha(q, k, torch.randn(3, 5, 24), q_mask, k_mask)
    with pytest.raises(ValueError):
        bad_k = torch.randn(2, 6, 24)
        bad_k_mask = _pad_mask(bad_k.shape[0], mha.num_heads, bad_k.shape[1], bad_k.device)
        mha(q, bad_k, v, q_mask, bad_k_mask)


# =======================
# Forward: happy path + masks + zero-length
# =======================


@pytest.mark.parametrize("B,S,D,H", [(1, 1, 16, 4), (2, 5, 32, 4)])
def test_forward_shapes_and_device_dtype(B, S, D, H):
    device = (
        torch.device(f"cuda:{torch.cuda.current_device()}")
        if torch.cuda.is_available()
        else torch.device("cpu")
    )
    mha = MultiHeadAttention(D, H, 0.1).to(device)
    q = _rand(B, S, D, device=device)
    k = _rand(B, S, D, device=device)
    v = _rand(B, S, D, device=device)
    mask = _pad_mask(B, H, S, device)

    out = mha(q, k, v, mask, mask)
    assert out.shape == (B, S, D)
    assert out.device == device
    assert out.dtype == q.dtype


def test_boolean_mask_blocks_positions():
    B, S, D, H = 2, 6, 24, 3
    mha = MultiHeadAttention(D, H, 0.0)  # no dropout for determinism
    q = _rand(B, S, D)
    k = _rand(B, S, D)
    v = _rand(B, S, D)

    q_mask = _pad_mask(B, H, S, q.device)
    k_mask = _pad_mask(B, H, S, q.device)
    k_mask[..., -2:] = True

    out1 = mha(q, k, v, q_mask, q_mask)
    out2 = mha(q, k, v, q_mask, k_mask)
    assert not torch.allclose(out1, out2)


def test_causal_mask_blocks_future_positions():
    B, S, D, H = 2, 5, 32, 4
    mha = MultiHeadAttention(D, H, 0.0)
    q, k, v = _rand(B, S, D), _rand(B, S, D), _rand(B, S, D)
    base_mask = _pad_mask(B, H, S, q.device)
    causal = torch.ones(B, H, S, S, dtype=torch.bool, device=q.device).triu(1)

    out_free = mha(q, k, v, base_mask, base_mask, None)
    out_causal = mha(q, k, v, base_mask, base_mask, causal)
    assert not torch.allclose(out_free, out_causal)


# =======================
# Gradients smoke
# =======================


def test_gradients_flow():
    mha = MultiHeadAttention(32, 4, 0.1)
    q = _rand(
        2,
        5,
        32,
        dtype=torch.float32,
        device="cpu",
    )
    k = _rand(
        2,
        5,
        32,
        dtype=torch.float32,
        device="cpu",
    )
    v = _rand(
        2,
        5,
        32,
        dtype=torch.float32,
        device="cpu",
    )
    q.requires_grad_(True)
    k.requires_grad_(True)
    v.requires_grad_(True)
    mask = _pad_mask(2, 4, 5, q.device)

    out = mha(q, k, v, mask, mask)
    loss = out.pow(2).mean()
    loss.backward()

    for t in (q, k, v):
        assert t.grad is not None
        assert torch.isfinite(t.grad).all()