"""Print every number the lesson computes (numbers from a cited page are marked "source").

    python3 show_numbers.py            uses the saved copy of the model config
    python3 show_numbers.py --fetch    downloads the live config and compares it
"""
import sys

from context_window import attention_pairs, count_tokens, tokens_read, total_tokens, truncate
from kv_memory import CONFIG_URL, fetch_config, kv_bytes, kv_bytes_at_max_context, kv_bytes_per_token, load_config

KIB, MIB, GIB = 1024, 1024**2, 1024**3

cfg = load_config()
if "--fetch" in sys.argv:
    print("live config equals saved copy:", fetch_config() == cfg)

print("== Qwen3-8B config:", CONFIG_URL)
for key in ("num_hidden_layers", "num_key_value_heads", "num_attention_heads", "head_dim",
            "torch_dtype", "max_position_embeddings"):
    print(f"  {key} = {cfg[key]}")
per_token = kv_bytes_per_token(cfg)
print(f"bytes per token = {per_token:,} ({per_token / KIB:g} KiB)")
for tokens in (1_000, 4_096, 32_768, 40_960):
    b = kv_bytes(cfg, tokens)
    print(f"  {tokens:>6,} tokens: {b:>14,} bytes = {b / MIB:,.3f} MiB = {b / GIB:.4f} GiB")
at_max = kv_bytes_at_max_context(cfg)
print(f"at max context: {at_max:,} bytes = {at_max / GIB:g} GiB = {at_max / 1e9:.2f} GB (decimal), "
      f"'about 6 gigabytes' in the opening")
mha = per_token * cfg["num_attention_heads"] // cfg["num_key_value_heads"]
print(f"if every query head had its own key/value pair: {mha:,} bytes per token "
      f"(x{cfg['num_attention_heads'] // cfg['num_key_value_heads']}), "
      f"{mha * cfg['max_position_embeddings'] / GIB:g} GiB at max context")
print(f"weights, about 8.2e9 parameters x 2 bytes = {8.2e9 * 2 / 1e9:.1f} GB (decimal; 8.2B is from the model card)")

print("\n== Window sizes (source: Google token guide lists 1,048,576; Qwen3-8B card says 32,768 natively)")
print(f"1,048,576 tokens at 60 to 80 words per 100 tokens (source ratio): "
      f"{1_048_576 * 0.6:,.0f} to {1_048_576 * 0.8:,.0f} words, "
      f"rounded to the nearest 10,000: {round(1_048_576 * 0.6, -4):,.0f} to {round(1_048_576 * 0.8, -4):,.0f}")

print("\n== Worked example (toy tokens = words)")
chat = [
    {"role": "system", "content": "Reply in French."},
    {"role": "user", "content": "My cat is Biscuit."},
    {"role": "assistant", "content": "Joli nom."},
    {"role": "user", "content": "I like tea."},
    {"role": "assistant", "content": "Moi aussi."},
    {"role": "user", "content": "What is my cat called?"},
]
print("sizes:", [count_tokens(m["content"]) for m in chat], "total", sum(count_tokens(m["content"]) for m in chat))
for pin in (True, False):
    kept = truncate(chat, 12, keep_system=pin)
    print(f"limit 12, keep_system={pin}: kept {[m['content'] for m in kept]}, {total_tokens(kept)} tokens")

print("\n== Re-reading history (every message 50 tokens)")
for n in (10, 20):
    text = 2 * n * 50
    read = tokens_read(n, 50)
    print(f"{n} replies: text written {text:,}, tokens read {read:,}, ratio {read / text:g}")
for k in range(1, 6):
    print(f"  reply {k}: {2 * k - 1} messages, reads {tokens_read(k, 50) - tokens_read(k - 1, 50):,}, "
          f"running total {tokens_read(k, 50):,}")
for n in range(1, 11):
    print(f"  reply {n}: reads {tokens_read(n, 50) - tokens_read(n - 1, 50):,} uncapped, "
          f"{tokens_read(n, 50, window=450) - tokens_read(n - 1, 50, window=450):,} with a 450-token window")
print(f"10 replies, window 450: {tokens_read(10, 50, window=450):,} read vs {tokens_read(10, 50):,} uncapped")
w40 = tokens_read(40, 50, window=450)
print(f"40 replies, window 450: {w40:,} read vs {tokens_read(40, 50):,} uncapped; "
      f"50 x (1+3+5+7+9) = {50 * (1 + 3 + 5 + 7 + 9):,}, then 450 x 35 = {450 * 35:,}")
print(f"window 400 holds 8 messages of 50 but sends 7: replies 5 to 10 read "
      f"{tokens_read(5, 50, window=400) - tokens_read(4, 50, window=400)} each; "
      f"10 replies read {tokens_read(10, 50, window=400):,}")

print("\n== Figure: seven messages and a 190-token window")
figure = [30, 20, 40, 25, 35, 30, 20]  # token counts, drawn at 2 px per token
print(f"messages {figure}; messages 3 to 7 = {sum(figure[2:])} tokens, plus 40 tokens of room for the reply "
      f"= a {sum(figure[2:]) + 40}-token window")


print("\n== Attention pairs per head and layer")
for n in (1_000, 2_000, 40_960):
    print(f"  {n:,} tokens: {attention_pairs(n):,} pairs; ratio to half the size: "
          f"{attention_pairs(n) / attention_pairs(n // 2):.3f}")

print("\n== Proof it works")
print(f"tokens_read(3 replies, 10 tokens each) = {tokens_read(3, 10)}")

print("\n== Check yourself")
q3 = [("system", 2), ("user", 3), ("assistant", 1), ("user", 6), ("assistant", 4), ("user", 3)]
q3_chat = [{"role": r, "content": " ".join(["w"] * n)} for r, n in q3]
kept = truncate(q3_chat, 10)
print(f"Q3 sizes {[n for _, n in q3]}, limit 10: kept sizes {[count_tokens(m['content']) for m in kept]} "
      f"roles {[m['role'] for m in kept]}, {total_tokens(kept)} tokens; pinned 2, budget 8; newest 3 leaves 5; "
      f"4-token reply leaves 1; 6 does not fit")
q4 = {"num_hidden_layers": 28, "num_key_value_heads": 4, "head_dim": 128, "torch_dtype": "float16",
      "hidden_size": 2048, "num_attention_heads": 16}
b4 = kv_bytes_per_token(q4)
print(f"Q4 per token {b4:,} bytes ({b4 / KIB:g} KiB); 8,192 tokens {kv_bytes(q4, 8_192):,} bytes "
      f"= {kv_bytes(q4, 8_192) / MIB:g} MiB")
print(f"Q5 12 replies of 40-token messages: {tokens_read(12, 40):,} read (40 x 12^2 = {40 * 12**2:,}); "
      f"window 200: {tokens_read(12, 40, window=200):,} = 40 + 120 + 200 + 9 x 200 = "
      f"{40 + 120 + 200 + 9 * 200:,}; first three replies read {tokens_read(3, 40, window=200):,}; "
      f"other 9 replies read 9 x 200 = {9 * 200:,}")