"""
two_phases.py: the Word Factory from Part 1, upgraded so the Typewriter can
type for SEVERAL users in one trip. We time:
  1. reading (prefill) prompts of different lengths, for one user
  2. typing (decode) one token each for 1, 4, 16 and 64 users at once
  3. what happens to everyone's typing when a long prompt arrives

GPT-2-small shape (12 floors, width 768), random weights, plain NumPy, CPU.
Tables are stored as (out, in) and applied as W @ x, which this CPU's math
library handles better for a handful of tokens.
"""
import time
import numpy as np

D, LAYERS, HEADS, D_FF, VOCAB = 768, 12, 12, 3072, 256
HD = D // HEADS

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)))  # not x**3 (slow!)

class Engine:
    def __init__(self, seed=0):
        rng = np.random.default_rng(seed)
        w = lambda *s: (rng.standard_normal(s) * 0.02).astype(np.float32)
        self.emb = w(VOCAB, D)
        self.layers = [dict(qkv=w(3 * D, D), out=w(D, D), up=w(D_FF, D), down=w(D, D_FF))
                       for _ in range(LAYERS)]
        self.head = w(VOCAB, D)

    def new_user(self, max_len):
        # every user gets their OWN notebook (KV cache)
        return dict(K=np.zeros((LAYERS, HEADS, max_len, HD), np.float32),
                    V=np.zeros((LAYERS, HEADS, max_len, HD), np.float32), n=0)

    def trip(self, x, segments):
        """One trip through all 12 floors.
        x: (D, T) numbers for T tokens; segments: [(user, n_tokens), ...]
        The big tables are shared by every token in the trip.
        Each user's attention reads only that user's notebook."""
        for i, p in enumerate(self.layers):
            q, k, v = np.split(p["qkv"] @ norm(x), 3)            # shared tables
            outs, col = [], 0
            for u, t in segments:                               # per-user notebook
                s, e = u["n"], u["n"] + t
                qi = q[:, col:col + t].reshape(HEADS, HD, t).transpose(0, 2, 1)  # (H, t, HD)
                u["K"][i, :, s:e] = k[:, col:col + t].reshape(HEADS, HD, t).transpose(0, 2, 1)
                u["V"][i, :, s:e] = v[:, col:col + t].reshape(HEADS, HD, t).transpose(0, 2, 1)
                Kc, Vc = u["K"][i, :, :e], u["V"][i, :, :e]      # (H, e, HD): read ALL notes
                att = qi @ Kc.transpose(0, 2, 1) / HD ** 0.5     # (H, t, e)
                if t > 1:                                        # no peeking ahead
                    att = np.where(np.triu(np.ones((t, e), bool), k=s + 1), -1e9, att)
                att = np.exp(att - att.max(-1, keepdims=True))
                att /= att.sum(-1, keepdims=True)
                outs.append((att @ Vc).transpose(0, 2, 1).reshape(D, t))
                col += t
            x = x + p["out"] @ np.concatenate(outs, 1)
            x = x + p["down"] @ gelu(p["up"] @ norm(x))
        for u, t in segments:
            u["n"] += t
        assert x.dtype == np.float32
        return x

    def prefill(self, user, ids):
        return self.trip(self.emb[ids].T.copy(), [(user, len(ids))])

    def decode_step(self, users, tok=7):
        x = np.repeat(self.emb[tok][:, None], len(users), 1)   # one token per user
        return self.trip(x, [(u, 1) for u in users])

def timed(fn, reps=5):
    ts = []
    for _ in range(reps):
        t = time.perf_counter(); fn(); ts.append(time.perf_counter() - t)
    return sorted(ts)[len(ts) // 2]

if __name__ == "__main__":
    eng = Engine()
    rng = np.random.default_rng(1)
    params = sum(a.size for p in eng.layers for a in p.values())
    print(f"tables in the 12 floors: {params:,} numbers ({params*4/1e6:.0f} MB)")

    print("\n1) READING (prefill): one user, prompts of different lengths")
    print(f"{'prompt':>7}{'time ms':>10}{'ms/token':>10}{'GFLOP/s':>9}")
    eng.prefill(eng.new_user(64), rng.integers(0, VOCAB, 32))          # warm-up
    for T in [1, 16, 64, 256, 1024]:
        ids = rng.integers(0, VOCAB, T)
        def run():
            eng.prefill(eng.new_user(T), ids)
        s = timed(run)
        print(f"{T:>7}{s*1e3:>10.1f}{s/T*1e3:>10.2f}{2*params*T/s/1e9:>9.1f}")

    print("\n2) TYPING (decode): one new token for each of B users, each with 512 notes")
    print(f"{'users':>7}{'step ms':>10}{'tokens/s':>10}{'GFLOP/s':>9}{'MB read':>9}")
    for B in [1, 4, 16, 64]:
        users = [eng.new_user(1024) for _ in range(B)]
        for u in users:
            u["n"] = 512                                     # pretend 512 earlier tokens
        def run():
            for u in users: u["n"] = 512                     # same spot every rep
            eng.decode_step(users)
        s = timed(run, reps=7)
        notes_mb = B * 512 * 2 * LAYERS * D * 4 / 1e6         # every user's notebook
        print(f"{B:>7}{s*1e3:>10.1f}{B/s:>10.1f}{2*params*B/s/1e9:>9.1f}{params*4/1e6 + notes_mb:>9.0f}")

    print("\n3) A 1,024-token prompt arrives while 8 users are typing")
    users = [eng.new_user(2048) for _ in range(8)]
    for u in users: u["n"] = 512
    def step():
        for u in users: u["n"] = 512
        eng.decode_step(users)
    base = timed(step, reps=5)
    ids = rng.integers(0, VOCAB, 1024)
    def whole():
        eng.prefill(eng.new_user(1024), ids)
    stall = timed(whole)
    def chunk():
        nu = eng.new_user(1024); nu["n"] = 896            # the last (hardest) chunk
        eng.prefill(nu, ids[:128])
    c = timed(chunk)
    print(f"normal gap between words for the 8 typers : {base*1e3:8.1f} ms")
    print(f"gap if the engine reads all 1,024 at once : {(base+stall)*1e3:8.1f} ms (one huge freeze)")
    print(f"gap if it reads 128 tokens between steps  : {(base+c)*1e3:8.1f} ms (8 small bumps, worst one)")