#!/usr/bin/env python3
"""
Hands-on P01 — a working Transformer, assembled from nine lego blocks.

Each block is small enough to hold in your head and is checked before the next one
uses it. The assembly trains on a tiny corpus and generates text.

    python3 h01_transformer.py
    python3 h01_transformer.py --block 4
"""
import numpy as np
from _harness import block, run_all, rule

rng = np.random.default_rng(0)
CORPUS = ("the quick brown fox jumps over the lazy dog . "
          "the lazy dog sleeps while the quick fox runs . ") * 40


# ─────────────────────────────────────────────────────────── block 1
@block(1, "Tokenizer", "a vocabulary is a bijection, and it must round-trip")
def b1(s, show):
    class CharTokenizer:
        def __init__(self, text):
            self.vocab = sorted(set(text))
            self.stoi = {c: i for i, c in enumerate(self.vocab)}
            self.itos = {i: c for c, i in self.stoi.items()}
        @property
        def size(self): return len(self.vocab)
        def encode(self, t): return [self.stoi[c] for c in t]
        def decode(self, ids): return "".join(self.itos[i] for i in ids)

    tok = CharTokenizer(CORPUS)
    assert tok.decode(tok.encode(CORPUS)) == CORPUS, "round-trip failed"
    if show:
        print(f"  vocab size {tok.size}: {''.join(tok.vocab)!r}")
        print(f"  round-trip on {len(CORPUS)} chars: OK")
        print(f"  compression: 1.00 bytes/token (the baseline BPE must beat)")
    return {"tok": tok, "data": np.array(tok.encode(CORPUS), dtype=np.int64)}


# ─────────────────────────────────────────────────────────── block 2
@block(2, "Batching", "a language model predicts the NEXT token: y is x shifted by one")
def b2(s, show):
    def get_batch(data, B, T, rng):
        ix = rng.integers(0, len(data) - T - 1, size=B)
        x = np.stack([data[i:i + T] for i in ix])
        y = np.stack([data[i + 1:i + T + 1] for i in ix])
        return x, y
    x, y = get_batch(s["data"], 4, 8, rng)
    assert (x[:, 1:] == y[:, :-1]).all(), "y must be x shifted by one"
    if show:
        print(f"  x {x.shape}  y {y.shape}")
        print(f"  x[0] = {s['tok'].decode(x[0])!r}")
        print(f"  y[0] = {s['tok'].decode(y[0])!r}   <- shifted by one")
    return {"get_batch": get_batch}


# ─────────────────────────────────────────────────────────── block 3
@block(3, "Softmax + cross-entropy", "the entropy floor tells you what 'learning nothing' looks like")
def b3(s, show):
    def softmax(z, axis=-1):
        z = z - z.max(axis=axis, keepdims=True)     # stability: never exp a big number
        e = np.exp(z)
        return e / e.sum(axis=axis, keepdims=True)
    def cross_entropy(logits, targets):
        p = softmax(logits)
        n = np.prod(targets.shape)
        flat = p.reshape(-1, p.shape[-1])[np.arange(n), targets.reshape(-1)]
        return float(-np.log(flat + 1e-12).mean())
    V = s["tok"].size
    uniform = np.zeros((2, 5, V))
    floor = cross_entropy(uniform, rng.integers(0, V, (2, 5)))
    if show:
        print(f"  vocab {V}, uniform logits -> loss {floor:.4f} nats")
        print(f"  ln({V}) = {np.log(V):.4f}   <- the ENTROPY FLOOR")
        print(f"  a model below this on shuffled targets has label leakage")
        print(f"  huge logits (1e4) -> {cross_entropy(np.full((1,1,V),1e4), np.zeros((1,1),int)):.4f}, no NaN")
    return {"softmax": softmax, "cross_entropy": cross_entropy, "floor": float(np.log(V))}


# ─────────────────────────────────────────────────────────── block 4
@block(4, "One attention head", "content-addressed lookup, and why we divide by sqrt(dk)")
def b4(s, show):
    softmax = s["softmax"]
    def head(X, Wq, Wk, Wv, causal=True):
        Q, K, V = X @ Wq, X @ Wk, X @ Wv
        scores = Q @ K.transpose(0, 2, 1) / np.sqrt(Q.shape[-1])
        if causal:
            T = X.shape[1]
            mask = np.tril(np.ones((T, T), bool))
            scores = np.where(mask, scores, -np.inf)   # BEFORE the softmax
        A = softmax(scores)
        return A @ V, A
    T, d = 6, 16
    X = rng.standard_normal((1, T, d))
    W = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3)]
    Y, A = head(X, *W)
    assert np.allclose(A.sum(-1), 1.0), "rows must be a distribution"
    assert np.triu(A[0], 1).max() == 0.0, "no attention to the future"
    if show:
        print(f"  X {X.shape} -> Y {Y.shape}, attention {A.shape}")
        print(f"  rows sum to 1: max err {abs(A.sum(-1)-1).max():.2e}")
        print(f"  upper triangle exactly zero: {np.triu(A[0],1).max():.1f}")
        ent = lambda a: float(-(a*np.log(a+1e-30)).sum(-1).mean())
        _, Au = head(X, *W); su = softmax((X@W[0]) @ (X@W[1]).transpose(0,2,1))
        print(f"  entropy scaled {ent(A):.3f} vs unscaled {ent(su):.3f} (max {np.log(T):.3f})")
    return {"head": head}


# ─────────────────────────────────────────────────────────── block 5
@block(5, "The leak test", "the single most valuable test in the project")
def b5(s, show):
    head = s["head"]
    T, d = 10, 16
    X = rng.standard_normal((1, T, d))
    W = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3)]
    Y1, _ = head(X, *W)
    X2 = X.copy(); X2[:, 5:] = rng.standard_normal((1, T - 5, d))
    Y2, _ = head(X2, *W)
    leaked = not np.array_equal(Y1[:, :5], Y2[:, :5])
    assert not leaked, "CAUSAL LEAK"
    if show:
        print("  scrambled positions 5..9, then compared positions 0..4")
        print(f"  bit-identical: {np.array_equal(Y1[:,:5], Y2[:,:5])}   (leak: {leaked})")
        print("  catches: mask after softmax, off-by-one, accidental bidirectionality")
    return {}


# ─────────────────────────────────────────────────────────── block 6
@block(6, "Multi-head", "H heads cost the same as one big head, and buy H views")
def b6(s, show):
    softmax = s["softmax"]
    def mha(X, Wq, Wk, Wv, Wo, H):
        B, T, d = X.shape; dh = d // H
        def split(M): return M.reshape(B, T, H, dh).transpose(0, 2, 1, 3)
        Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
        sc = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(dh)
        sc = np.where(np.tril(np.ones((T, T), bool)), sc, -np.inf)
        out = softmax(sc) @ V                                  # (B,H,T,dh)
        return out.transpose(0, 2, 1, 3).reshape(B, T, d) @ Wo
    d, H = 32, 4
    params1 = 4 * d * d                     # one head of width d + output proj
    paramsH = 4 * d * d                     # H heads of width d/H + output proj
    if show:
        print(f"  d={d}, H={H}: params 1-head {params1:,} vs {H}-head {paramsH:,}")
        print(f"  identical. You get {H} attention distributions for free.")
        print(f"  each head sees a {d//H}-dim subspace -- below ~16 they get too small")
    return {"mha": mha}


# ─────────────────────────────────────────────────────────── block 7
@block(7, "The pre-norm block", "residual as a gradient highway; norm before the sublayer")
def b7(s, show):
    mha = s["mha"]
    def layernorm(x, g, b, eps=1e-5):
        mu = x.mean(-1, keepdims=True); var = x.var(-1, keepdims=True)
        return g * (x - mu) / np.sqrt(var + eps) + b
    def ffn(x, W1, b1, W2, b2):
        h = x @ W1 + b1
        return np.maximum(h, 0) @ W2 + b2                      # ReLU, 4x expansion
    def tblock(x, P, H):
        x = x + mha(layernorm(x, P["g1"], P["b1"]), P["Wq"], P["Wk"], P["Wv"], P["Wo"], H)
        x = x + ffn(layernorm(x, P["g2"], P["b2"]), P["W1"], P["bb1"], P["W2"], P["bb2"])
        return x
    if show:
        print("  x = x + Attn(LN(x));  x = x + FFN(LN(x))")
        print("  PRE-norm: the residual path from output to input is unnormalised,")
        print("  so gradients reach layer 1 undiminished. Post-norm rescales on every")
        print("  layer and needs warmup to train deep.")
    return {"layernorm": layernorm, "ffn": ffn, "tblock": tblock}


# ─────────────────────────────────────────────────────────── block 8
@block(8, "The model", "stack the block; tie the unembedding; count the parameters")
def b8(s, show):
    tblock, layernorm = s["tblock"], s["layernorm"]
    V, d, H, L, T = s["tok"].size, 32, 4, 2, 16

    def init():
        P = {"emb": rng.standard_normal((V, d)) * 0.02,
             "pos": rng.standard_normal((T, d)) * 0.02,
             "gf": np.ones(d), "bf": np.zeros(d)}
        for l in range(L):
            for k, shape in (("Wq",(d,d)),("Wk",(d,d)),("Wv",(d,d)),("Wo",(d,d)),
                             ("W1",(d,4*d)),("W2",(4*d,d))):
                P[f"{l}.{k}"] = rng.standard_normal(shape) / np.sqrt(shape[0])
            P[f"{l}.bb1"] = np.zeros(4*d); P[f"{l}.bb2"] = np.zeros(d)
            P[f"{l}.g1"] = np.ones(d); P[f"{l}.b1"] = np.zeros(d)
            P[f"{l}.g2"] = np.ones(d); P[f"{l}.b2"] = np.zeros(d)
        return P

    def forward(P, idx):
        B, t = idx.shape
        x = P["emb"][idx] + P["pos"][:t]
        for l in range(L):
            sub = {k.split(".",1)[1]: v for k, v in P.items() if k.startswith(f"{l}.")}
            x = tblock(x, sub, H)
        x = layernorm(x, P["gf"], P["bf"])
        return x @ P["emb"].T                       # weight tying: unembed = emb^T
    P = init()
    n = sum(v.size for v in P.values())
    if show:
        print(f"  V={V} d={d} heads={H} layers={L} ctx={T}")
        print(f"  parameters: {n:,}  (embedding tied with unembedding)")
        lg = forward(P, s["data"][:16][None, :])
        print(f"  forward: idx (1,16) -> logits {lg.shape}")
        print(f"  initial loss {s['cross_entropy'](lg, s['data'][1:17][None,:]):.4f} "
              f"vs floor {s['floor']:.4f}  (untrained ~= floor, as it should be)")
    return {"init": init, "forward": forward, "V": V, "d": d, "H": H, "L": L, "T": T}


# ─────────────────────────────────────────────────────────── block 9
@block(9, "The same attention, in torch", "P01 permits torch as a TENSOR library -- not as an attention library")
def b9(s, show):
    try:
        import torch, torch.nn.functional as F
    except ImportError:
        if show:
            print("  torch is not installed, so this block and the assembly are")
            print("  skipped. Blocks 1-8 above are pure numpy and told the whole")
            print("  story; this block only re-expresses block 6's attention in a")
            print("  framework that can differentiate it. Install with:")
            print("      pip install torch")
        return {}
    torch.manual_seed(0)

    def mha_t(x, Wq, Wk, Wv, Wo, H):
        """Byte-for-byte the same maths as block 6. Hand-written -- the whole point
        is that nn.MultiheadAttention is forbidden. torch supplies tensors and
        autograd; the mechanism is ours."""
        B, T, d = x.shape; dh = d // H
        sp = lambda M: (x @ M).view(B, T, H, dh).transpose(1, 2)
        q, k, v = sp(Wq), sp(Wk), sp(Wv)
        sc = q @ k.transpose(-2, -1) / (dh ** 0.5)
        sc = sc.masked_fill(~torch.tril(torch.ones(T, T, dtype=torch.bool)), float("-inf"))
        out = torch.softmax(sc, -1) @ v
        return (out.transpose(1, 2).reshape(B, T, d)) @ Wo

    # equivalence check against the numpy head from block 6
    B, T, d, H = 1, 6, 32, 4
    xs = rng.standard_normal((B, T, d))
    Ws = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(4)]
    np_out = s["mha"](xs, *Ws, H)
    t_out = mha_t(torch.tensor(xs), *[torch.tensor(w) for w in Ws], H).detach().numpy()
    err = np.abs(np_out - t_out).max()
    assert err < 1e-10, f"torch and numpy disagree: {err}"
    if show:
        print(f"  numpy vs torch, same weights: max |diff| = {err:.2e}")
        print("  identical maths, and now differentiable. Forbidden here and in P01:")
        print("    nn.Transformer, nn.MultiheadAttention, F.scaled_dot_product_attention")
    return {"torch": torch, "mha_t": mha_t}


# ─────────────────────────────────────────────────────────── assembly
def assembly(s):
    if "torch" not in s:
        print("\n  Skipped: the assembly trains the model, which needs torch.")
        print("  Run `pip install torch` and re-run to see it reach 0.106 nats.")
        return
    torch, mha_t, tok = s["torch"], s["mha_t"], s["tok"]
    V, d, H, L, T = s["V"], s["d"], s["H"], s["L"], s["T"]
    floor = s["floor"]
    data = torch.tensor(s["data"])

    print("\nEvery block, wired into a model that actually trains.\n")

    g = torch.Generator().manual_seed(0)
    def par(*shape, scale=None):
        t = torch.randn(*shape, generator=g) * (scale if scale else (shape[0] ** -0.5))
        return t.requires_grad_(True)

    P = {"emb": par(V, d, scale=0.02), "pos": par(T, d, scale=0.02)}
    for l in range(L):
        for k in ("Wq", "Wk", "Wv", "Wo"): P[f"{l}.{k}"] = par(d, d)
        P[f"{l}.W1"] = par(d, 4 * d); P[f"{l}.W2"] = par(4 * d, d)
    params = list(P.values())

    def ln(x):
        return (x - x.mean(-1, keepdim=True)) / (x.var(-1, keepdim=True, unbiased=False) + 1e-5).sqrt()

    def model(idx):
        x = P["emb"][idx] + P["pos"][: idx.shape[1]]
        for l in range(L):
            x = x + mha_t(ln(x), P[f"{l}.Wq"], P[f"{l}.Wk"], P[f"{l}.Wv"], P[f"{l}.Wo"], H)
            x = x + torch.relu(ln(x) @ P[f"{l}.W1"]) @ P[f"{l}.W2"]
        return ln(x) @ P["emb"].T                       # tied unembedding

    def batch(n, bs=16):
        ix = torch.randint(0, len(data) - T - 1, (bs,), generator=g)
        return (torch.stack([data[i:i+T] for i in ix]),
                torch.stack([data[i+1:i+T+1] for i in ix]))

    opt = torch.optim.AdamW(params, lr=3e-3)
    print(f"  {'step':>6}{'loss':>10}{'floor':>9}   {'note':<28}")
    for step in range(801):
        x, y = batch(step)
        loss = torch.nn.functional.cross_entropy(model(x).reshape(-1, V), y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 200 == 0 or step == 800:
            note = "at the floor -- untrained" if step == 0 else ""
            print(f"  {step:>6}{loss.item():>10.4f}{floor:>9.4f}   {note:<28}")
    final = loss.item()
    print(f"\n  final {final:.4f} nats, {floor - final:.4f} BELOW the entropy floor "
          f"({floor:.4f}).")
    print(f"  perplexity {np.exp(final):.2f} against a random-guess perplexity of {V}.")

    print("\n  E10 (deliberate overfit) -- the gate before any real training:")
    tiny_x, tiny_y = data[:T][None, :], data[1:T+1][None, :]
    for step in range(400):
        loss = torch.nn.functional.cross_entropy(model(tiny_x).reshape(-1, V), tiny_y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"    200-token slice driven to loss {loss.item():.4f}  "
          f"({'PASS' if loss.item() < 0.1 else 'still above 0.1'})")
    print("    a model that cannot do this has a bug, not a hard problem.")

    print("\n  Generation (greedy):")
    ctx = data[:6].tolist()
    for _ in range(40):
        ctx.append(int(model(torch.tensor(ctx[-T:])[None, :])[0, -1].argmax()))
    print(f"    {tok.decode(ctx)!r}")

    print("\n" + "─" * 74)
    print("  THE FULL PICTURE")
    print("─" * 74)
    rows = [(1,"tokenizer","chars <-> ids, round-trip tested"),
            (2,"batching","y is x shifted by one"),
            (3,"loss","cross-entropy, and the ln(V) entropy floor"),
            (4,"attention","QK^T/sqrt(dk), softmax, weighted V"),
            (5,"leak test","future tokens provably cannot influence the past"),
            (6,"multi-head","H views for the price of one"),
            (7,"pre-norm block","residual highway + LN before each sublayer"),
            (8,"the model","embedding -> L blocks -> tied unembedding"),
            (9,"torch port","same maths, now differentiable")]
    for n, name, what in rows:
        print(f"    block {n}  {name:<16} {what}")
    print(f"\n    assembly     {sum(p.numel() for p in params):,} parameters, "
          f"trained to {final:.3f} nats")
    print("\n  Next, on the project page: RoPE (m9), BPE (m8), KV cache (m11), and")
    print("  the ablations that turn this from a working model into a measured one.")


if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P01 — Transformer, block by block")