"""Tests on the serve-quant recording: Lab 12's runs of Qwen3 on one laptop (Apple M4 Max, 464 GB/s
measured), the 4-bit Qwen3-4B's changed answers on 300 made-up support emails, and its recorded
answers on the long-pile test. Run: pytest test_quant_practice.py"""
from quant_practice import ceiling, compare, recall, weights_gb

BW = 464  # GB/s, a matrix x vector read of 8 GB (script 07)


def test_ceiling():
    assert round(ceiling(BW, 16.4)) == 28  # Qwen3-8B at 16 bits; measured 25 tokens/s
    assert round(ceiling(BW, 8.0)) == 58  # Qwen3-4B; measured 45
    assert 25 / ceiling(BW, 16.4) > 0.85  # close to the ceiling: decoding mostly waits for memory


def test_weights_gb():
    # Qwen3-4B has about 4.02 billion numbers: 8.0, 4.3, 2.3 and 1.8 GB recorded
    assert [round(weights_gb(4.02, b), 1) for b in (16, 8, 4, 3)] == [8.0, 4.3, 2.3, 1.8]
    # Qwen3-8B, about 8.19 billion: 16.4, 8.7, 4.6, 3.6 GB
    assert [round(weights_gb(8.19, b), 1) for b in (16, 8, 4, 3)] == [16.4, 8.7, 4.6, 3.6]
    # the fine-tuned 0.6B, about 0.596 billion: 1.19, 0.63, 0.34, 0.26 GB
    assert [round(weights_gb(0.596, b), 2) for b in (16, 8, 4, 3)] == [1.19, 0.63, 0.34, 0.26]
    # 4 bits is just under a third of 16, not a quarter: the per-group rulers cost half a bit each
    assert 0.25 < weights_gb(1, 4) / weights_gb(1, 16) < 1 / 3


def test_compare_counts_changed_answers():
    # The 4-bit Qwen3-4B against 16-bit on 300 emails, rebuilt from the recorded counts (not the answers
    # themselves): 176 right both times, 22 right -> wrong, 3 wrong -> right, 25 wrong -> a different
    # wrong, 74 wrong the same way.
    gold = ["ok"] * 300
    ref = ["ok"] * 176 + ["ok"] * 22 + ["a"] * 3 + ["a"] * 25 + ["a"] * 74
    new = ["ok"] * 176 + ["a"] * 22 + ["ok"] * 3 + ["b"] * 25 + ["a"] * 74
    c = compare(gold, ref, new)
    assert (c["ref_right"], c["new_right"]) == (198, 179)  # 66% -> 60%: 6 points
    assert (c["changed"], c["r2w"], c["w2r"], c["w2w"]) == (50, 22, 3, 25)
    assert round(c["same"], 2) == 0.83  # 1 answer in 6 changed
    assert compare(["x"], ["x"], ["x"])["same"] == 1.0


# The long-pile test at about 4,000 tokens (50-54 emails each): Qwen3-4B's recorded answers.
NAMES_4K = ["Aiko", "Marcus", "Fatima", "Chloe", "Chloe", "Fatima", "Hannah", "Chloe", "Chloe", "Marcus"]
AT_16 = ["Aiko", "Marcus", "Fatima", "Oliver", "Chloe", "Fatima", "Lucia", "Chloe", "Chloe", "Marcus"]
AT_4 = ["Hannah", "Ben", "Ravi", "Oliver", "Aiko", "Elena", "Mei", "Ravi", "Mei", "Marcus"]


def test_recall_long_pile():
    assert recall(AT_16, NAMES_4K) == 8
    assert recall(AT_4, NAMES_4K) == 1  # lost 6 points on the short test, 7 of 10 here
    assert recall([" fatima. ", "Fatima"], ["Fatima", "fatima"]) == 2


def test_recall_grades_like_the_lab():
    # At 3 bits and ~1,500 tokens the model wrapped a right name in markdown; the lab counts it wrong.
    assert recall(["**Marcus**"], ["Marcus"]) == 0
    assert recall(["Marcus."], ["Marcus"]) == 1
    assert recall(["A-7783 is mentioned in the"], ["Chloe"]) == 0