"""Guards for memory_calc.py: the parameter count, the byte sizes and the fit test."""
import ast
import re
from pathlib import Path

import pytest

import memory_calc as mc

# Totals shown on each model page as "Model size" (Hugging Face, read 2026-10-19).
PUBLISHED = {"1b": 1_235_814_400, "8b": 8_030_261_248, "70b": 70_553_706_496, "405b": 405_853_388_800}


@pytest.mark.parametrize("key", PUBLISHED)
def test_counted_parameters_equal_the_published_total(key):
    """Guards the shape formula (a missed norm, key/value width, or tied output matrix is off by millions)."""
    assert mc.param_count(mc.load_config(key)) == PUBLISHED[key]


# The invented shape of the hand-worked table in the lesson.
TINY = {"num_hidden_layers": 2, "hidden_size": 8, "intermediate_size": 16, "num_attention_heads": 4,
        "num_key_value_heads": 2, "vocab_size": 10}


def test_tiny_model_matches_the_lessons_hand_table():
    """Guards the lesson's table: each row was added up on paper, and the code must land on the same bytes."""
    assert mc.head_dim(TINY) == 2
    assert mc.param_count(TINY) == 2 * (2 * 8 * 8 + 2 * 8 * 4 + 3 * 8 * 16 + 2 * 8) + 10 * 8 + 10 * 8 + 8 == 1_352
    assert mc.weight_bytes(TINY, 16) == 2_704
    assert mc.kv_bytes_per_token(TINY) == 32
    assert mc.kv_cache_bytes(TINY, 100, 3) == 9_600
    assert mc.activation_bytes(TINY, tokens=10) == 10 * (3 * 16 + 4 * 8) * 2 == 1_600


def test_tied_output_matrix_is_counted_once():
    """Guards counting a second vocab x hidden matrix for the 1B model, which reuses its token table."""
    cfg = mc.load_config("1b")
    untied = dict(cfg, tie_word_embeddings=False)
    table = cfg["vocab_size"] * cfg["hidden_size"]
    assert mc.param_count(untied) - mc.param_count(cfg) == table
    assert mc.weight_bytes(cfg, 8) == 2 * 262_668_288 + 973_146_112 == 1_498_482_688
    assert mc.weight_bytes(untied, 8) - mc.weight_bytes(cfg, 8) == 2 * table   # the extra matrix is 16-bit


def test_weights_follow_bits_per_weight_by_hand():
    """Guards the embeddings: 1,050,673,152 weights (token table and output matrix) stay at 16 bits, the other 6,979,588,096 use `bits`."""
    cfg = mc.load_config("8b")
    assert all(isinstance(mc.weight_bytes(cfg, bits), int) for bits in mc.BITS)   # a byte count is whole: `//`, not `/`
    assert mc.weight_bytes(cfg, 16) == 16_060_522_496 == 2 * mc.param_count(cfg)
    assert mc.weight_bytes(cfg, 8) == 2 * 1_050_673_152 + 6_979_588_096 == 9_080_934_400
    assert mc.weight_bytes(cfg, 4) == 2 * 1_050_673_152 + 6_979_588_096 // 2 == 5_591_140_352
    with pytest.raises(ValueError):
        mc.weight_bytes(cfg, 3)


# Sizes of the .safetensors files of published 8-bit and 4-bit copies (Hugging Face API, read 2026-10-19).
FILES = [("8b", 8, 9_081_287_016), ("8b", 4, 5_727_938_576), ("8b", 4, 5_735_719_888), ("70b", 4, 39_767_996_256)]


@pytest.mark.parametrize("key,bits,size", FILES)
def test_weights_sit_just_under_published_quantized_files(key, bits, size):
    """Guards counting the embeddings at the low precision, which comes 11% to 30% short of these files;
    the scale values a real 4-bit file adds are the few percent left over."""
    gap = (size - mc.weight_bytes(mc.load_config(key), bits)) / size
    assert 0 <= gap < 0.04


@pytest.mark.parametrize("key,per_token", [("1b", 32_768), ("8b", 131_072), ("70b", 327_680), ("405b", 516_096)])
def test_cache_bytes_per_token_by_hand(key, per_token):
    """2 x layers x key/value heads x head size x 2 bytes, worked out on paper for each model."""
    assert mc.kv_bytes_per_token(mc.load_config(key)) == per_token


def test_cache_grows_with_tokens_and_chats_and_8b_at_8192_is_one_gib():
    cfg = mc.load_config("8b")
    assert mc.kv_cache_bytes(cfg, 8192, 1) == mc.GIB
    assert mc.kv_cache_bytes(cfg, 8192, 8) == 8 * mc.GIB
    assert mc.kv_cache_bytes(cfg, 16384, 1) == 2 * mc.GIB
    assert mc.kv_cache_bytes(cfg, 8192, 0) == 0


def test_explicit_head_size_wins_over_hidden_divided_by_heads():
    """Guards a model whose head size is not hidden / heads (the cache would be sized wrong)."""
    cfg = {"hidden_size": 1024, "num_attention_heads": 16, "head_dim": 128, "num_key_value_heads": 8, "num_hidden_layers": 28}
    assert mc.head_dim(cfg) == 128
    assert mc.kv_bytes_per_token(cfg) == 2 * 28 * 8 * 128 * 2
    assert mc.head_dim({"hidden_size": 1024, "num_attention_heads": 16}) == 64


def test_working_memory_by_hand():
    # 2048 tokens x (3 x 14336 + 4 x 4096) values x 2 bytes
    assert mc.activation_bytes(mc.load_config("8b")) == 2048 * (43008 + 16384) * 2 == 243_269_632


def test_total_is_the_sum_of_the_four_parts():
    m = mc.memory_need("8b", 16, 8192, 4)
    assert m["total"] == m["weights"] + m["kv_cache"] + m["activations"] + m["overhead"]
    assert m["kv_cache"] == 4 * mc.GIB and m["overhead"] == mc.GIB


def test_fit_uses_the_usable_share_and_the_edge_is_inclusive():
    """Guards testing against the whole card: a server claims 92% by default."""
    limit = mc.usable_bytes(24)
    assert limit == int(24 * mc.GIB * 0.92)
    assert mc.fits(limit, 24) and not mc.fits(limit + 1, 24)
    assert mc.fits(mc.memory_need("8b", 16, 8192, 4)["total"], 24)
    assert not mc.fits(mc.memory_need("8b", 16, 8192, 8)["total"], 24)


def test_cards_needed_rounds_up():
    limit = mc.usable_bytes(80)
    assert mc.cards_needed(limit, 80) == 1 and mc.cards_needed(limit + 1, 80) == 2
    assert mc.cards_needed(mc.memory_need("70b", 16, 8192, 1)["total"], 80) == 2


def test_chats_that_fit_leaves_the_fixed_parts_alone():
    # 24 GiB card: 22.08 usable - 14.96 weights - 0.23 working - 1 fixed leaves room for five 1 GiB chats, not six.
    assert mc.chats_that_fit("8b", 16, 8192, 24) == 5
    assert mc.fits(mc.memory_need("8b", 16, 8192, 5)["total"], 24)
    assert not mc.fits(mc.memory_need("8b", 16, 8192, 6)["total"], 24)
    assert mc.chats_that_fit("70b", 16, 8192, 80) == 0   # the weights alone do not fit
    assert mc.chats_that_fit("70b", 8, 8192, 80) == 1    # 73.60 usable - 67.67 weights - 0.45 working - 1 fixed = 4.48; a chat is 2.50


def _without_docstring(source, node):
    """The lines of one function, minus its docstring."""
    lines = source.splitlines()[node.lineno - 1:node.end_lineno]
    first = node.body[0]
    if isinstance(first, ast.Expr) and isinstance(first.value, ast.Constant) and isinstance(first.value.value, str):
        del lines[first.lineno - node.lineno:first.end_lineno - node.lineno + 1]
    return "\n".join(lines)


def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
    """The lesson shows `weight_bytes` and `memory_need` with the docstrings left out. If the code changes and the
    page does not, the reader copies a function that was never tested."""
    lesson = (Path(__file__).parent.parent / "lesson.md").read_text()
    source = Path(mc.__file__).read_text()
    in_source = {n.name: _without_docstring(source, n) for n in ast.parse(source).body if isinstance(n, ast.FunctionDef)}
    blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
    shown = [node for block in blocks for node in ast.parse(block).body]
    assert {node.name for node in shown} == {"weight_bytes", "memory_need"}
    for block in blocks:
        for node in ast.parse(block).body:
            assert _without_docstring(block, node) == in_source[node.name], f"{node.name} differs from memory_calc.py"