"""
forward_pass.py: follow ONE token through a GPT-2-small-shaped model and
count the work (FLOPs) and the time spent in every part.

Weights are random: the words are nonsense, but the amount of work is exactly
what a trained model of this shape does. FLOPs: 1 multiply + 1 add = 2 FLOPs,
so a (rows x cols) matrix times a vector costs 2 * rows * cols FLOPs.
"""
import sys, time
from collections import defaultdict
import numpy as np

D, LAYERS, HEADS, D_FF, VOCAB = 768, 12, 12, 3072, 50257     # GPT-2 small
HD = D // HEADS

flops = defaultdict(float)    # part name -> FLOPs for ONE token
secs  = defaultdict(float)    # part name -> seconds (summed over all steps)

def mm(part, x, W):
    """Vector x matrix, with the bookkeeping."""
    t = time.perf_counter()
    y = x @ W
    secs[part] += time.perf_counter() - t
    flops[part] += 2 * W.shape[0] * W.shape[1]
    return y

def timed(part, fn, *args, n_flops=0):
    t = time.perf_counter()
    y = fn(*args)
    secs[part] += time.perf_counter() - t
    flops[part] += n_flops
    return y

def norm(x):
    return (x - x.mean()) / np.sqrt(x.var() + 1e-5)

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

class Model:
    def __init__(self, n_ctx, seed=0):
        rng = np.random.default_rng(seed)
        w = lambda *s: (rng.standard_normal(s) * 0.02).astype(np.float32)
        self.tok_emb = w(VOCAB, D)                       # the "dictionary"
        self.layers = [dict(qkv=w(D, 3 * D), out=w(D, D),
                            up=w(D, D_FF), down=w(D_FF, D))
                       for _ in range(LAYERS)]
        self.lm_head = self.tok_emb.T.copy()             # GPT-2 reuses it
        # The notebook (KV cache), already holding n_ctx earlier tokens.
        self.K = w(LAYERS, HEADS, n_ctx, HD)
        self.V = w(LAYERS, HEADS, n_ctx, HD)
        self.n_ctx = n_ctx

    def params(self):
        layer = sum(a.size for a in self.layers[0].values())
        return {"per layer": layer, "all layers": layer * LAYERS,
                "dictionary": self.tok_emb.size}

    def one_token(self, tok):
        n = self.n_ctx
        x = timed("1. look up the token", lambda: self.tok_emb[tok].copy())
        for p in self.layers:
            # Look-back room (attention)
            h = timed("other (norms, softmax, adds)", norm, x, n_flops=5 * D)
            q, k, v = np.split(mm("2. attention: make Q, K, V", h, p["qkv"]), 3)
            q = q.reshape(HEADS, HD)
            # compare the new token with every note in the notebook
            K, V = self.K[self.layer_i], self.V[self.layer_i]
            s = timed("3. attention: read the notebook",
                      # HD ** 0.5, not np.sqrt(HD): a NumPy float64 here would
                      # silently turn all the math after it into slow 64-bit math
                      lambda: np.einsum("hd,hnd->hn", q, K) / HD ** 0.5,
                      n_flops=2 * HEADS * n * HD)
            a = timed("other (norms, softmax, adds)",
                      lambda: np.exp(s - s.max(-1, keepdims=True)), n_flops=3 * HEADS * n)
            a /= a.sum(-1, keepdims=True)
            mix = timed("3. attention: read the notebook",
                        lambda: np.einsum("hn,hnd->hd", a, V).reshape(D),
                        n_flops=2 * HEADS * n * HD)
            x = x + mm("4. attention: output", mix, p["out"])
            # Thinking room (MLP)
            h = timed("other (norms, softmax, adds)", norm, x, n_flops=5 * D)
            up = mm("5. thinking: expand 768 -> 3072", h, p["up"])
            up = timed("other (norms, softmax, adds)", gelu, up, n_flops=8 * D_FF)
            x = x + mm("6. thinking: shrink 3072 -> 768", up, p["down"])
            self.layer_i += 1
        assert x.dtype == np.float32, "something switched us to 64-bit math"
        h = timed("other (norms, softmax, adds)", norm, x, n_flops=5 * D)
        logits = mm("7. score every word in the dictionary", h, self.lm_head)
        return int(logits.argmax())

    def run(self, steps):
        tok = 464                                       # any starting token
        for _ in range(steps):
            self.layer_i = 0
            tok = self.one_token(tok)

def report(n_ctx, steps=30, repeats=3):
    m = Model(n_ctx)
    m.run(3)                                            # warm-up
    runs = []
    for _ in range(repeats):                            # time it 3 times...
        flops.clear(); secs.clear()
        m.run(steps)
        runs.append((sum(secs.values()), dict(flops), dict(secs)))
    _, f_all, s_all = sorted(runs, key=lambda r: r[0])[repeats // 2]   # ...keep the middle run
    total_f = sum(f_all.values()) / steps
    total_s = sum(s_all.values()) / steps
    print(f"\nnotebook holds {n_ctx:,} tokens | one token = "
          f"{total_f/1e6:,.1f} million FLOPs, {total_s*1e3:.1f} ms")
    print(f"{'part':40s}{'MFLOPs':>9s}{'work %':>8s}{'ms':>8s}{'time %':>8s}")
    for k in sorted(f_all):
        f, t = f_all[k] / steps, s_all[k] / steps
        print(f"{k:40s}{f/1e6:9.2f}{100*f/total_f:7.1f}%{t*1e3:8.2f}{100*t/total_s:7.1f}%")
    return m, total_f, total_s

if __name__ == "__main__":
    ctxs = [int(c) for c in sys.argv[1:]] or [128, 1024, 8192]
    for i, n in enumerate(ctxs):
        m, f, s = report(n)
        if i == 0:
            p = m.params()
            print(f"\nparameters: {p['all layers']:,} in the 12 layers "
                  f"({p['per layer']:,} each) + {p['dictionary']:,} in the dictionary")
            print(f"total {p['all layers'] + p['dictionary']:,} -> 2 x params = "
                  f"{2*(p['all layers'] + p['dictionary'])/1e6:,.1f} million FLOPs")
            print(f"speed while typing one token: {f/s/1e9:.2f} GFLOP/s")