import torch
from transformers import Qwen2Config

from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV


def tiny_expert(non_causal_attention: bool = False) -> KVCachedDiffusionExpert:
    vlm_config = Qwen2Config(
        vocab_size=32,
        hidden_size=64,
        intermediate_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        max_position_embeddings=128,
    )
    architecture = ExpertArchitecture(
        hidden_size=32,
        intermediate_size=64,
        num_layers=2,
        num_attention_heads=2,
        num_key_value_heads=2,
        head_dim=16,
        action_projection_hidden_size=16,
        action_fourier_features=4,
        non_causal_attention=non_causal_attention,
    )
    return KVCachedDiffusionExpert(vlm_config, architecture)


def test_expert_consumes_compatible_past_kv() -> None:
    expert = tiny_expert()
    prefix = PrefixKV(
        keys=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
        values=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
    )
    output = expert(torch.randn(1, 3, 2), torch.rand(1, 1, 1), prefix)
    assert output.shape == (1, 3, 2)


def test_expert_rejects_incompatible_kv_heads() -> None:
    expert = tiny_expert()
    prefix = PrefixKV(
        keys=tuple(torch.randn(1, 1, 5, 16) for _ in range(2)),
        values=tuple(torch.randn(1, 1, 5, 16) for _ in range(2)),
    )
    try:
        expert.validate_prefix(prefix)
    except ValueError as error:
        assert "KV-head mismatch" in str(error)
    else:
        raise AssertionError("Expected incompatible KV heads to be rejected")


def test_non_causal_expert_allows_future_action_attention() -> None:
    torch.manual_seed(3)
    expert = tiny_expert(non_causal_attention=True).eval()
    prefix = PrefixKV(
        keys=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
        values=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
    )
    first = torch.randn(1, 3, 2)
    changed_future = first.clone()
    changed_future[:, -1] += 0.137
    timestep = torch.rand(1, 1, 1)
    with torch.no_grad():
        first_output = expert(first, timestep, prefix)
        second_output = expert(changed_future, timestep, prefix)
    assert not torch.allclose(first_output[:, 0], second_output[:, 0])