import numpy as np

from kv_sizes import (LLAMA3_8B, OPT_13B, decode_both_ways, kv_bytes, kv_bytes_per_token,
                      pairs_with_cache, pairs_without_cache)


def test_formula_matches_public_configurations():
    # Llama 3 8B: 2 x 32 layers x 8 kv heads x 128 x 2 bytes = 128 KiB per token
    assert kv_bytes_per_token(LLAMA3_8B) == 131072
    assert kv_bytes(LLAMA3_8B, 8192) == 1024**3          # exactly 1 GiB at the full 8,192-token context
    # The PagedAttention paper computes OPT-13B as 2 x 5120 x 40 x 2 bytes = 800 KB per token
    assert kv_bytes_per_token(OPT_13B) == 2 * 5120 * 40 * 2 == 819200


def test_cache_changes_the_cost_not_the_answer():
    """A cache that reads the wrong rows or is off by one would change the outputs."""
    n = 60
    without, with_cache, pairs_a, pairs_b = decode_both_ways(n)
    assert np.allclose(without, with_cache)
    assert (pairs_a, pairs_b) == (pairs_without_cache(n), pairs_with_cache(n)) == (n * (n + 1) // 2, n)