"""Tensor-shape contracts for single-head and multi-head attention."""

import pytest
import torch

from attention import (
    MultiHeadSelfAttention,
    SelfAttention,
    causal_mask,
    combine_causal_and_padding,
    scaled_dot_product_attention,
)


def test_single_head_shapes():
    batch, seq_len, d_model, d_k, d_v = 2, 4, 8, 4, 6
    x = torch.randn(batch, seq_len, d_model)
    module = SelfAttention(d_model=d_model, d_k=d_k, d_v=d_v)

    output, weights = module(x)

    assert module.W_Q.shape == (d_model, d_k)
    assert module.W_K.shape == (d_model, d_k)
    assert module.W_V.shape == (d_model, d_v)
    assert output.shape == (batch, seq_len, d_v)
    assert weights.shape == (batch, seq_len, seq_len)
    assert torch.isfinite(output).all()
    assert torch.isfinite(weights).all()


def test_scores_and_weights_have_query_by_key_shape():
    batch, seq_len, d_head = 3, 7, 5
    q = torch.randn(batch, seq_len, d_head)
    k = torch.randn(batch, seq_len, d_head)
    v = torch.randn(batch, seq_len, d_head)

    output, weights = scaled_dot_product_attention(q, k, v, return_weights=True)

    assert weights.shape == (batch, seq_len, seq_len)
    assert output.shape == (batch, seq_len, d_head)
    # Every query row is a distribution over keys.
    torch.testing.assert_close(
        weights.sum(dim=-1), torch.ones(batch, seq_len), atol=1e-6, rtol=0
    )


def test_multi_head_split_merge_shapes():
    batch, seq_len, num_heads, d_head = 2, 5, 4, 8
    d_model = num_heads * d_head
    x = torch.randn(batch, seq_len, d_model)
    module = MultiHeadSelfAttention(d_model=d_model, num_heads=num_heads)

    qkv = module.in_proj(x)
    q, k, v = qkv.chunk(3, dim=-1)
    split_q = q.view(batch, seq_len, num_heads, d_head).transpose(1, 2)

    output, weights = module(x, mask=causal_mask(seq_len))

    assert qkv.shape == (batch, seq_len, 3 * d_model)
    assert q.shape == k.shape == v.shape == (batch, seq_len, d_model)
    assert split_q.shape == (batch, num_heads, seq_len, d_head)
    assert weights.shape == (batch, num_heads, seq_len, seq_len)
    assert output.shape == (batch, seq_len, d_model)

    merged = output.new_empty(batch, seq_len, d_model)
    assert merged.shape == output.shape


def test_multi_head_rejects_indivisible_width():
    with pytest.raises(ValueError, match="divisible"):
        MultiHeadSelfAttention(d_model=10, num_heads=3)


def test_combined_mask_broadcasts_to_scores():
    seq_len, batch, heads = 4, 2, 3
    causal = causal_mask(seq_len)
    key_padding = torch.tensor(
        [[True, True, False, False], [True, True, True, False]]
    )

    combined = combine_causal_and_padding(causal, key_padding)

    assert causal.shape == (seq_len, seq_len)
    assert combined.shape == (batch, 1, seq_len, seq_len)

    scores = torch.randn(batch, heads, seq_len, seq_len)
    masked = scores.masked_fill(~combined, float("-inf"))
    assert masked.shape == scores.shape

    for b in range(batch):
        for t in range(seq_len):
            for j in range(seq_len):
                expected = bool(causal[t, j]) and bool(key_padding[b, j])
                assert bool(combined[b, 0, t, j]) is expected


def test_mask_shape_errors_are_explicit():
    with pytest.raises(ValueError, match="causal must be"):
        combine_causal_and_padding(torch.ones(3, 4, dtype=torch.bool),
                                   torch.ones(1, 4, dtype=torch.bool))

    q = torch.randn(1, 4, 8)
    with pytest.raises(ValueError, match="sequence length"):
        scaled_dot_product_attention(q, torch.randn(1, 5, 8), q)

    with pytest.raises(ValueError, match="head width"):
        scaled_dot_product_attention(q, torch.randn(1, 4, 6), q)