import torch

from grimboost.config import DrafterConfig
from grimboost.model.transformer import GrimBoostLM


def tiny_config() -> DrafterConfig:
    return DrafterConfig(
        vocab_size=256,
        hidden_size=64,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=16,
        intermediate_size=128,
        max_position_embeddings=64,
        rope_theta=10_000.0,
    )


def make_model() -> GrimBoostLM:
    torch.manual_seed(1)
    return GrimBoostLM(tiny_config()).eval()


def test_incremental_equals_full():
    model = make_model()
    ids = torch.randint(0, 256, (1, 12))
    with torch.no_grad():
        full = model(ids)
        cache = model.new_cache(max_len=64, dtype=torch.float32)
        _ = model(ids[:, :-1], past_key_values=cache, use_cache=True)
        step = model(ids[:, -1:], past_key_values=cache, use_cache=True)
    assert cache.length == 12
    assert torch.allclose(step[:, -1], full[:, -1], atol=1e-5)


def test_multi_token_block_after_cache():
    model = make_model()
    ids = torch.randint(0, 256, (1, 14))
    with torch.no_grad():
        full = model(ids)
        cache = model.new_cache(max_len=64, dtype=torch.float32)
        _ = model(ids[:, :8], past_key_values=cache, use_cache=True)
        block = model(ids[:, 8:], past_key_values=cache, use_cache=True)
    assert cache.length == 14
    assert torch.allclose(block, full[:, 8:], atol=1e-5)


def test_newest_suffix_rollback_is_exact():
    model = make_model()
    ids = torch.randint(0, 256, (1, 14))
    alt = torch.randint(0, 256, (1, 3))
    with torch.no_grad():
        cache = model.new_cache(max_len=64, dtype=torch.float32)
        _ = model(ids[:, :11], past_key_values=cache, use_cache=True)
        _ = model(ids[:, 11:], past_key_values=cache, use_cache=True)
        assert cache.length == 14
        cache.rollback(3)
        assert cache.length == 11
        block = model(alt, past_key_values=cache, use_cache=True)
        reference = model(torch.cat([ids[:, :11], alt], dim=1))
    assert torch.allclose(block[:, -1], reference[:, -1], atol=1e-5)


def test_rebuild_with_positions_from_zero():
    model = make_model()
    ids = torch.randint(0, 256, (1, 10))
    with torch.no_grad():
        full = model(ids)
        cache = model.new_cache(max_len=64, dtype=torch.float32)
        rebuilt = model(ids, past_key_values=cache, use_cache=True)
    assert cache.length == 10
    assert torch.allclose(rebuilt, full, atol=1e-5)


def test_left_crop_is_not_equivalent():
    """Documents why the engine never left-crops a full-attention cache."""
    model = make_model()
    ids = torch.randint(0, 256, (1, 12))
    with torch.no_grad():
        full_suffix = model(ids[:, 5:], position_ids=torch.arange(5, 12).unsqueeze(0))
        cache = model.new_cache(max_len=64, dtype=torch.float32)
        _ = model(ids[:, :-1], past_key_values=cache, use_cache=True)
        for layer in range(model.config.num_hidden_layers):
            cache.k[layer] = cache.k[layer][:, :, 5:, :]
            cache.v[layer] = cache.v[layer][:, :, 5:, :]
        cache.length -= 5
        cropped = model(ids[:, -1:], past_key_values=cache, use_cache=True,
                        position_ids=torch.tensor([[11]]))
    assert not torch.allclose(cropped[:, -1], full_suffix[:, -1], atol=1e-4)