"""
kv_cache.py: build the Word Factory's notebook (the KV cache) by hand.
  1. Same answer? Generate words WITHOUT a notebook (re-read everything every
     step) and WITH one, and check both give the same words.
  2. Time per new word, with and without the notebook, as the text grows.
  3. How big the notebook is, measured from the arrays themselves.
  4. A trick to shrink it: let several heads share one notebook (GQA),
     and time how long reading a long notebook takes each way.
GPT-2-small shape (12 floors, width 768, 12 heads), random weights, NumPy, CPU.
"""
import time
import numpy as np

D, LAYERS, HEADS, D_FF, VOCAB = 768, 12, 12, 3072, 256
HD = D // HEADS
rng = np.random.default_rng(0)
w = lambda *s: (rng.standard_normal(s) * 0.02).astype(np.float32)
EMB, POS = w(VOCAB, D), w(4096, D)
L = [dict(qkv=w(3 * D, D), out=w(D, D), up=w(D_FF, D), down=w(D, D_FF)) for _ in range(LAYERS)]
HEAD = w(VOCAB, D)

def norm(x):                                   # x: (D, T)
    return (x - x.mean(0)) / np.sqrt(x.var(0) + 1e-5)

def gelu(x):
    return 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))

def softmax(a):
    a = np.exp(a - a.max(-1, keepdims=True))
    return a / a.sum(-1, keepdims=True)

def heads(m):                                  # (D, T) -> (H, T, HD)
    return m.reshape(HEADS, HD, -1).transpose(0, 2, 1)

def run(tokens, start, cache=None):
    """Push `tokens` (sitting at positions start..) through all 12 floors.
    With a cache: write their K, V into it and read everything stored so far.
    Without one: the tokens can only look at each other."""
    t = len(tokens)
    x = (EMB[tokens] + POS[start:start + t]).T  # (D, t)
    for i, p in enumerate(L):
        q, k, v = (heads(m) for m in np.split(p["qkv"] @ norm(x), 3))
        if cache is not None:
            cache["K"][i, :, start:start + t] = k
            cache["V"][i, :, start:start + t] = v
            k, v = cache["K"][i, :, :start + t], cache["V"][i, :, :start + t]
        att = q @ k.transpose(0, 2, 1) / HD ** 0.5                 # (H, t, all)
        mask = np.triu(np.ones((t, k.shape[1]), bool), k=k.shape[1] - t + 1)
        att = softmax(np.where(mask, -1e9, att))
        x = x + p["out"] @ (att @ v).transpose(0, 2, 1).reshape(D, t)
        x = x + p["down"] @ gelu(p["up"] @ norm(x))
    return HEAD @ norm(x)[:, -1]               # scores for the next word

def new_cache(n):
    return dict(K=np.zeros((LAYERS, HEADS, n, HD), np.float32),
                V=np.zeros((LAYERS, HEADS, n, HD), np.float32))

# ---- 1) Same answer with and without the notebook ----------------------
prompt = list(rng.integers(0, VOCAB, 128))
NEW = 32

def generate_no_cache():
    toks, times, scores = list(prompt), [], []
    for _ in range(NEW):
        t = time.perf_counter()
        logits = run(toks, 0)                  # re-read EVERYTHING, every step
        times.append(time.perf_counter() - t)
        scores.append(logits); toks.append(int(logits.argmax()))
    return toks[len(prompt):], times, scores

def generate_with_cache():
    c = new_cache(len(prompt) + NEW)
    logits = run(prompt, 0, c)                 # read the question once
    toks, times, scores = [int(logits.argmax())], [], [logits]
    for i in range(NEW - 1):
        t = time.perf_counter()
        logits = run([toks[-1]], len(prompt) + i, c)   # only the newest word
        times.append(time.perf_counter() - t)
        scores.append(logits); toks.append(int(logits.argmax()))
    return toks, times, scores

a, t_no, s_no = generate_no_cache()
b, t_yes, s_yes = generate_with_cache()
gap = max(np.abs(x - y).max() for x, y in zip(s_no, s_yes))
print("1) same words with and without the notebook:", a == b)
print(f"   biggest difference in any score: {gap:.7f}")
print("   first 10 words:", a[:10])

# ---- 2) Time per new word ---------------------------------------------
print("\n2) time per new word (ms)")
print("   text length   no notebook   with notebook")
for i in [1, 8, 16, 31]:
    print(f"   {len(prompt) + i:>11} {t_no[i] * 1e3:13.1f} {t_yes[i - 1] * 1e3:15.1f}")
print(f"   total for {NEW} words: {sum(t_no):.2f} s without, {sum(t_yes):.2f} s with")

# ---- 3) How big is the notebook? ---------------------------------------
c = new_cache(1)
per_token = c["K"].nbytes + c["V"].nbytes
print(f"\n3) notebook per token: {per_token:,} bytes = {per_token / 1024:.0f} KB")
for n in [1_000, 8_000, 32_000]:
    print(f"   {n:>6,} tokens: {per_token * n / 1e6:7.0f} MB   (the model itself is 340 MB)")

# ---- 4) Share notebooks between heads (GQA) -----------------------------
T = 8_192
print(f"\n4) reading a {T:,}-token notebook for ONE new word, all 12 floors")
for kv_heads in [12, 3, 1]:
    K = np.random.rand(LAYERS, kv_heads, T, HD).astype(np.float32)
    V = np.random.rand(LAYERS, kv_heads, T, HD).astype(np.float32)
    q = np.random.rand(LAYERS, kv_heads, HEADS // kv_heads, HD).astype(np.float32)
    def read():
        for i in range(LAYERS):                # each group of query heads
            att = softmax(q[i] @ K[i].transpose(0, 2, 1))   # shares one K, V
            att @ V[i]
    read(); ts = []
    for _ in range(5):
        t = time.perf_counter(); read(); ts.append(time.perf_counter() - t)
    mb = (K.nbytes + V.nbytes) / 1e6
    print(f"   {HEADS} query heads, {kv_heads:>2} notebook heads: {mb:5.0f} MB read, "
          f"{np.median(ts) * 1e3:5.1f} ms")
    del K, V