#!/usr/bin/env python3
"""
W1 — Attention in 40 lines, and the two tests that matter.  (~45 min)

Miniature of P01. Build single-head causal attention with numpy, prove the causal
mask does not leak, prove the scale keeps softmax responsive, and find the sequence
length where the quadratic term overtakes the feed-forward.
"""
import numpy as np

rng = np.random.default_rng(0)


def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)        # stability: never exp a big number
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)


def attention(X, Wq, Wk, Wv, causal=True, scale=True):
    """X: (T,d). Returns (T,d)."""
    Q, K, V = X @ Wq, X @ Wk, X @ Wv               # (T,dk) each
    scores = Q @ K.T                               # (T,T)  <- the quadratic term
    if scale:
        scores = scores / np.sqrt(Q.shape[-1])
    if causal:
        T = X.shape[0]
        # -inf BEFORE the softmax, so masked positions get exactly zero weight.
        # Applied after, they would get a small nonzero weight and the model would
        # quietly cheat.
        scores = np.where(np.tril(np.ones((T, T), bool)), scores, -np.inf)
    A = softmax(scores)
    return A @ V, A


T, d = 12, 32
X = rng.standard_normal((T, d))
Wq, Wk, Wv = (rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3))

print("=" * 68)
print("TEST 1 — attention weights form a distribution")
_, A = attention(X, Wq, Wk, Wv)
print(f"  every row sums to 1: max |sum-1| = {np.abs(A.sum(1) - 1).max():.2e}")
print(f"  strictly lower-triangular support: upper triangle max = {np.triu(A,1).max():.2e}")

print()
print("TEST 2 — the causal mask does not leak")
Y1, _ = attention(X, Wq, Wk, Wv)
X2 = X.copy()
X2[6:] = rng.standard_normal((T - 6, d))           # scramble everything from row 6 on
Y2, _ = attention(X2, Wq, Wk, Wv)
same = np.array_equal(Y1[:6], Y2[:6])
print(f"  rows 0..5 bit-identical after scrambling rows 6..11: {same}")
print(f"  rows 6..11 changed:                                  {not np.allclose(Y1[6:], Y2[6:])}")
assert same, "CAUSAL LEAK"

print()
print("TEST 3 — what the 1/sqrt(dk) scale buys")
# Initialise W so that Q,K have unit-variance COMPONENTS. Without the 1/sqrt(dk)
# here, Q already has variance dk per component and the scores are astronomical
# whatever you divide by afterwards -- initialisation and scaling are two separate
# defences against the same failure, and the first draft of this script forgot one.
ent = lambda A: float(-(A * np.log(A + 1e-30)).sum(1).mean())
for dk in (16, 64, 256):
    Xk = rng.standard_normal((T, dk))
    W = [rng.standard_normal((dk, dk)) / np.sqrt(dk) for _ in range(3)]
    _, A_s = attention(Xk, *W, scale=True)
    _, A_u = attention(Xk, *W, scale=False)
    print(f"  dk={dk:>4}   mean row entropy: scaled {ent(A_s):.3f}   "
          f"unscaled {ent(A_u):.3f}   (max {np.log(T):.3f})")

print()
print("TEST 4 — where does the quadratic term overtake the FFN?")
print(f"  {'T':>6} {'attn GF':>9} {'ffn GF':>9} {'quad share':>11}")
H, dh = 12, 64
dm = H * dh
for Tq in (128, 512, 1024, 2048, 4096, 8192):
    qkv = 3 * 2 * Tq * dm * dm
    quad = 2 * 2 * 1 * H * Tq * Tq * dh            # scores + scores@V
    proj = 2 * Tq * dm * dm
    ffn = 16 * Tq * dm * dm
    attn = qkv + quad + proj
    print(f"  {Tq:>6} {attn/1e9:>8.2f}G {ffn/1e9:>8.2f}G {quad/(attn+ffn)*100:>10.1f}%")
print("\n  attention is ~18% of a GPT-2-small layer at its own context length of 1024.")
print("  'Attention is the bottleneck' is a claim about a context length.")