from kv_memory import (BYTES_PER_VALUE, kv_bytes, kv_bytes_at_max_context, kv_bytes_per_token,
                       load_config)


def test_qwen3_8b_matches_the_hand_computed_value():
    """2 x 36 layers x 8 kv heads x 128 x 2 bytes, worked out by hand: 147,456 bytes (144 KiB)."""
    cfg = load_config()
    assert cfg["num_hidden_layers"] == 36 and cfg["num_key_value_heads"] == 8
    assert cfg["head_dim"] == 128 and cfg["torch_dtype"] == "bfloat16"
    assert kv_bytes_per_token(cfg) == 2 * 36 * 8 * 128 * 2 == 147_456
    assert kv_bytes_per_token(cfg) == 144 * 1024


def test_memory_at_the_context_sizes_the_lesson_quotes():
    cfg = load_config()
    assert cfg["max_position_embeddings"] == 40_960
    assert kv_bytes_at_max_context(cfg) == 6_039_797_760 == int(5.625 * 1024**3)
    assert kv_bytes(cfg, 32_768) == 4_831_838_208 == int(4.5 * 1024**3)
    assert kv_bytes(cfg, 4_096) == 576 * 1024**2
    assert kv_bytes(cfg, 2 * 4_096) == 2 * kv_bytes(cfg, 4_096)      # linear in tokens


def test_uses_kv_heads_not_attention_heads():
    """The common slip: with 32 query heads instead of 8 key/value heads the answer is 4x too big."""
    cfg = load_config()
    assert cfg["num_attention_heads"] == 32
    assert kv_bytes_per_token(cfg) * 4 == 589_824
    assert kv_bytes_per_token({**cfg, "num_key_value_heads": 32}) == 589_824 != kv_bytes_per_token(cfg)


def test_explicit_head_dim_beats_hidden_size_division():
    """Made-up config: hidden 1024 / 16 heads = 64, but the file says head_dim 128. 2 x 2 x 1 x 128 x 2 = 1,024."""
    cfg = {"num_hidden_layers": 2, "num_key_value_heads": 1, "num_attention_heads": 16,
           "hidden_size": 1024, "head_dim": 128, "torch_dtype": "float16"}
    assert kv_bytes_per_token(cfg) == 1024


def test_missing_kv_heads_means_one_per_attention_head():
    """Made-up config without grouped-query attention: 2 x 1 x 4 x (8 / 4 = 2) x 4 bytes = 64."""
    cfg = {"num_hidden_layers": 1, "num_attention_heads": 4, "hidden_size": 8, "torch_dtype": "float32"}
    assert kv_bytes_per_token(cfg) == 64
    assert BYTES_PER_VALUE["float32"] == 4


def test_dtype_key_is_the_fallback_when_torch_dtype_is_missing():
    """Newer config files say `dtype`. 1 layer, 1 kv head, head size 8 / 1 = 8, 2 bytes: 2 x 8 x 2 = 32."""
    cfg = {"num_hidden_layers": 1, "num_key_value_heads": 1, "num_attention_heads": 1,
           "hidden_size": 8, "dtype": "bfloat16"}
    assert kv_bytes_per_token(cfg) == 32


def test_lessons_check_yourself_cache_question():
    """28 layers, 4 key/value heads, head size 128, 2 bytes per value: 57,344 bytes per token, 448 MiB at 8,192."""
    cfg = {"num_hidden_layers": 28, "num_key_value_heads": 4, "num_attention_heads": 16,
           "head_dim": 128, "torch_dtype": "float16"}
    assert kv_bytes_per_token(cfg) == 2 * 28 * 4 * 128 * 2 == 57_344
    assert kv_bytes(cfg, 8_192) == 469_762_048 == 448 * 1024**2