#!/usr/bin/env python3
"""
proofs.py — numerically verify every derivation the track relies on.

    python3 proofs.py            # run all, print PASS/FAIL per claim
    python3 proofs.py --verbose  # show the numbers behind each check

A proof on a page is an assertion until something checks it. Each function below
corresponds to one section of proofs.md and verifies the closed-form result against
a direct computation, a simulation, or a brute-force search.

This is the same discipline the track demands of you: the derivation is the claim, the
numerical check is the evidence, and shipping the first without the second is what
ai-policy.md rule 8 forbids.
"""

from __future__ import annotations

import math
import random
import sys
from itertools import combinations

VERBOSE = "--verbose" in sys.argv
RESULTS: list[tuple[str, bool, str]] = []


def check(name: str, ok: bool, detail: str = "") -> None:
    RESULTS.append((name, ok, detail))
    if VERBOSE or not ok:
        print(f"  [{'PASS' if ok else 'FAIL'}] {name}")
        if detail:
            print(f"         {detail}")


def close(a: float, b: float, rel: float = 1e-9) -> bool:
    return abs(a - b) <= rel * max(1.0, abs(a), abs(b))


# ---------------------------------------------------------------- P1 sqrt(d_k)
def p1_softmax_scale():
    """Var(q.k) = d_k for unit-variance components, so scores scale as sqrt(d_k)."""
    rng = random.Random(0)
    for d in (16, 64, 256, 1024):
        dots = []
        for _ in range(20000):
            q = [rng.gauss(0, 1) for _ in range(d)]
            k = [rng.gauss(0, 1) for _ in range(d)]
            dots.append(sum(a * b for a, b in zip(q, k)))
        mean = sum(dots) / len(dots)
        var = sum((x - mean) ** 2 for x in dots) / (len(dots) - 1)
        # variance should equal d; sampling error ~ d*sqrt(2/n)
        tol = d * 0.05
        check(f"P1 Var(q·k) = d_k at d={d}", abs(var - d) < tol,
              f"measured {var:.1f}, predicted {d}, |err| {abs(var-d):.1f} < {tol:.1f}")

    # and the consequence: unscaled softmax saturates
    def softmax(xs):
        m = max(xs); e = [math.exp(x - m) for x in xs]; s = sum(e)
        return [v / s for v in e]

    d = 64
    q = [rng.gauss(0, 1) for _ in range(d)]
    keys = [[rng.gauss(0, 1) for _ in range(d)] for _ in range(20)]
    raw = [sum(a * b for a, b in zip(q, k)) for k in keys]
    scaled = [x / math.sqrt(d) for x in raw]
    ent = lambda p: -sum(x * math.log(x) for x in p if x > 0)
    e_raw, e_scaled = ent(softmax(raw)), ent(softmax(scaled))
    check("P1 scaling raises softmax entropy (avoids saturation)", e_scaled > e_raw,
          f"entropy unscaled {e_raw:.3f} -> scaled {e_scaled:.3f} "
          f"(max possible {math.log(20):.3f})")


# ------------------------------------------------------------------- P2/P3 bloom
def p2_bloom_optimal_k():
    """k* = (m/n) ln2 minimises fpr = (1 - e^{-kn/m})^k, and fpr* = 0.6185^{m/n}."""
    for bpk in (4, 8, 10, 16, 20):
        f = lambda k: (1 - math.exp(-k / bpk)) ** k
        # brute-force the minimiser over a fine grid
        grid = [i / 100 for i in range(1, 4001)]
        kstar_num = min(grid, key=f)
        kstar_closed = bpk * math.log(2)
        check(f"P2 k* = (m/n)ln2 at {bpk} bits/key",
              abs(kstar_num - kstar_closed) < 0.02,
              f"numeric argmin {kstar_num:.3f}, closed form {kstar_closed:.3f}")

        fpr_closed = 0.6185 ** bpk
        check(f"P3 fpr* = 0.6185^(m/n) at {bpk} bits/key",
              abs(f(kstar_closed) - fpr_closed) / fpr_closed < 0.02,
              f"f(k*) {f(kstar_closed):.6f} vs 0.6185^{bpk} = {fpr_closed:.6f}")

    # at the optimum each bit is set with probability exactly 1/2
    bpk = 10; k = bpk * math.log(2)
    p_set = 1 - math.exp(-k / bpk)
    check("P2 at k* each bit is 1 with probability 1/2 (max entropy)",
          close(p_set, 0.5, 1e-12), f"P(bit=1) = {p_set:.12f}")


# ------------------------------------------------------------ P4 quorum
def p4_quorum_intersection():
    """Any two quorums of size Q from N intersect iff 2Q > N. Verified exhaustively."""
    for N in range(2, 10):
        for Q in range(1, N + 1):
            nodes = range(N)
            disjoint_exists = any(
                not (set(a) & set(b))
                for a in combinations(nodes, Q)
                for b in combinations(nodes, Q)
            )
            predicted_always_intersect = (2 * Q > N)
            check(f"P4 N={N} Q={Q}: intersect-always == (2Q>N)",
                  (not disjoint_exists) == predicted_always_intersect,
                  f"disjoint pair exists: {disjoint_exists}; 2Q>N: {predicted_always_intersect}")


# ------------------------------------------------------------ P5 Little's Law
def p5_littles_law():
    """L = lambda*W verified by simulating an open-loop arrival process."""
    rng = random.Random(1)
    lam = 500.0          # arrivals/sec
    W = 0.05             # 50 ms in system
    T = 400.0
    arrivals = []
    t = 0.0
    while t < T:
        t += rng.expovariate(lam)
        arrivals.append(t)
    # each item stays exactly W; integrate the count over time
    events = sorted([(a, +1) for a in arrivals] + [(a + W, -1) for a in arrivals])
    cur = 0; last = 0.0; area = 0.0
    for tm, d in events:
        area += cur * (tm - last); last = tm; cur += d
    L_measured = area / last
    L_predicted = lam * W
    check("P5 Little's Law L = lambda*W",
          abs(L_measured - L_predicted) / L_predicted < 0.05,
          f"measured L {L_measured:.2f}, predicted {L_predicted:.2f}")

    # the bytes-in-flight application from numbers.md section 2
    bw, lat = 57.5e9, 121e-9
    inflight = bw * lat
    check("P5 bytes in flight to sustain 57.5 GB/s at 121 ns",
          abs(inflight / 64 - 108.7) < 1.0,
          f"{inflight:.0f} bytes = {inflight/64:.1f} cache lines")


# ------------------------------------------------- P6 decode arithmetic intensity
def p6_decode_intensity():
    """I = 2*N*b / (N*bytes) = 2b/bytes: independent of model size N."""
    for bytes_per in (1, 2, 4):
        for b in (1, 8, 64, 512):
            vals = []
            for N in (1e8, 7e9, 7e10, 4e11):
                W = 2 * N * b
                Q = N * bytes_per
                vals.append(W / Q)
            check(f"P6 intensity independent of N (b={b}, {bytes_per}B)",
                  max(vals) - min(vals) < 1e-9,
                  f"I = {vals[0]:.1f} for every N; closed form 2b/bytes = {2*b/bytes_per:.1f}")
            assert close(vals[0], 2 * b / bytes_per)


# --------------------------------------------------------- P7 forward vs reverse AD
def p7_ad_modes():
    """Cost of a full Jacobian: forward O(n) passes, reverse O(m) passes."""
    # Build a random linear chain f: R^n -> R^m and count passes needed.
    rng = random.Random(3)
    n, m, L = 6, 2, 4
    dims = [n] + [5] * (L - 1) + [m]
    mats = [[[rng.gauss(0, 1) for _ in range(dims[i])] for _ in range(dims[i + 1])]
            for i in range(L)]

    def forward(x):
        v = x
        for M in mats:
            v = [sum(row[j] * v[j] for j in range(len(v))) for row in M]
        return v

    def jvp(tangent):           # one forward-mode pass -> one COLUMN of J
        v = tangent
        for M in mats:
            v = [sum(row[j] * v[j] for j in range(len(v))) for row in M]
        return v

    def vjp(cotangent):         # one reverse-mode pass -> one ROW of J
        v = cotangent
        for M in reversed(mats):
            v = [sum(v[i] * M[i][j] for i in range(len(M))) for j in range(len(M[0]))]
        return v

    # assemble J both ways and compare
    J_fwd = [[0.0] * n for _ in range(m)]
    for j in range(n):                       # n passes
        e = [1.0 if k == j else 0.0 for k in range(n)]
        col = jvp(e)
        for i in range(m):
            J_fwd[i][j] = col[i]
    J_rev = [[0.0] * n for _ in range(m)]
    for i in range(m):                       # m passes
        e = [1.0 if k == i else 0.0 for k in range(m)]
        J_rev[i] = vjp(e)

    err = max(abs(J_fwd[i][j] - J_rev[i][j]) for i in range(m) for j in range(n))
    check("P7 forward and reverse produce the same Jacobian", err < 1e-9,
          f"max |J_fwd - J_rev| = {err:.2e}; forward used {n} passes, reverse {m}")
    check("P7 reverse is cheaper exactly when m < n", (m < n) and (m < n),
          f"n={n} inputs, m={m} outputs -> reverse wins by {n/m:.1f}x")


# ------------------------------------------------------------ P8 matmul backward
def p8_matmul_backward():
    """For C = AB: dA = dC B^T and dB = A^T dC. Checked against finite differences."""
    rng = random.Random(4)
    M, K, N = 4, 3, 5
    A = [[rng.gauss(0, 1) for _ in range(K)] for _ in range(M)]
    B = [[rng.gauss(0, 1) for _ in range(N)] for _ in range(K)]
    G = [[rng.gauss(0, 1) for _ in range(N)] for _ in range(M)]   # upstream dL/dC

    def matmul(X, Y):
        return [[sum(X[i][k] * Y[k][j] for k in range(len(Y)))
                 for j in range(len(Y[0]))] for i in range(len(X))]

    def loss(A, B):
        C = matmul(A, B)
        return sum(C[i][j] * G[i][j] for i in range(M) for j in range(N))

    # analytic
    Bt = [[B[k][j] for k in range(K)] for j in range(N)]
    At = [[A[i][k] for i in range(M)] for k in range(K)]
    dA = matmul(G, Bt)
    dB = matmul(At, G)

    # finite differences
    h = 1e-6
    worst = 0.0
    for i in range(M):
        for k in range(K):
            Ap = [row[:] for row in A]; Ap[i][k] += h
            Am = [row[:] for row in A]; Am[i][k] -= h
            num = (loss(Ap, B) - loss(Am, B)) / (2 * h)
            worst = max(worst, abs(num - dA[i][k]) / (abs(num) + 1e-9))
    for k in range(K):
        for j in range(N):
            Bp = [row[:] for row in B]; Bp[k][j] += h
            Bm = [row[:] for row in B]; Bm[k][j] -= h
            num = (loss(A, Bp) - loss(A, Bm)) / (2 * h)
            worst = max(worst, abs(num - dB[k][j]) / (abs(num) + 1e-9))
    check("P8 dA = dC B^T and dB = A^T dC", worst < 1e-5,
          f"worst relative error vs finite differences: {worst:.2e}")


# ------------------------------------------------------------ P9 sample size
def p9_sample_size():
    """n = 2(z_a + z_b)^2 sigma^2 / delta^2, and halving delta quadruples n."""
    za, zb = 1.959963984540054, 0.8416212335729143
    bracket = (za + zb) ** 2
    check("P9 (z_0.975 + z_0.80)^2 = 7.849", abs(bracket - 7.8489) < 1e-3,
          f"{bracket:.6f} -- the folklore '16 sigma^2/delta^2' is 2*7.849 = {2*bracket:.2f}")

    def n_per_arm(sigma, delta):
        return math.ceil(2 * bracket * sigma ** 2 / delta ** 2)

    a, b, c = n_per_arm(0.5, 0.05), n_per_arm(0.5, 0.025), n_per_arm(0.5, 0.0125)
    check("P9 halving the MDE quadruples n",
          abs(b / a - 4) < 0.02 and abs(c / b - 4) < 0.02,
          f"n = {a}, {b}, {c}; ratios {b/a:.3f}, {c/b:.3f}")

    # empirical power check: does n_per_arm actually give 80% power?
    rng = random.Random(5)
    sigma, delta = 0.5, 0.1
    n = n_per_arm(sigma, delta)
    hits = 0; trials = 3000
    for _ in range(trials):
        A = [rng.gauss(0.0, sigma) for _ in range(n)]
        B = [rng.gauss(delta, sigma) for _ in range(n)]
        ma, mb = sum(A) / n, sum(B) / n
        va = sum((x - ma) ** 2 for x in A) / (n - 1)
        vb = sum((x - mb) ** 2 for x in B) / (n - 1)
        se = math.sqrt(va / n + vb / n)
        if abs(mb - ma) / se > za:
            hits += 1
    power = hits / trials
    check("P9 the formula delivers ~80% power", abs(power - 0.80) < 0.04,
          f"n={n} per arm, measured power {power:.3f} over {trials} trials")


# ------------------------------------------------------------ P10 peeking
def p10_peeking():
    """Repeated significance testing inflates the false-positive rate."""
    rng = random.Random(6)
    za = 1.959963984540054

    def fpr(nlooks, per_look=400, trials=3000):
        hits = 0
        for _ in range(trials):
            sa = ssa = sb = ssb = 0.0; n = 0
            for _ in range(nlooks):
                for _ in range(per_look):
                    x = rng.gauss(0, 1); y = rng.gauss(0, 1)
                    sa += x; ssa += x * x; sb += y; ssb += y * y
                n += per_look
                ma, mb = sa / n, sb / n
                va = (ssa - n * ma * ma) / (n - 1); vb = (ssb - n * mb * mb) / (n - 1)
                se = math.sqrt(va / n + vb / n)
                if se > 0 and abs(mb - ma) / se > za:
                    hits += 1; break
        return hits / trials

    f1, f5, f20 = fpr(1), fpr(5), fpr(20)
    check("P10 one look gives the nominal 5%", abs(f1 - 0.05) < 0.015,
          f"{f1*100:.2f}%")
    check("P10 five looks roughly triples it", f5 > 2 * f1,
          f"1 look {f1*100:.2f}% -> 5 looks {f5*100:.2f}%")
    check("P10 twenty looks exceeds 20%", f20 > 0.20,
          f"20 looks {f20*100:.2f}% -- a 'win' at p<0.05 is now 1-in-4 noise")


# ------------------------------------------------------------ P11 EMA half-life
def p11_ema_halflife():
    """h = ln(0.5)/ln(1-alpha): the lag at which an observation's weight halves."""
    for alpha in (0.02, 0.05, 0.1, 0.2, 0.5):
        h = math.log(0.5) / math.log(1 - alpha)
        w0 = alpha                       # weight of the most recent observation
        wh = alpha * (1 - alpha) ** h    # weight h steps back
        check(f"P11 half-life at alpha={alpha}", close(wh / w0, 0.5, 1e-9),
              f"h = {h:.2f} interactions; w(h)/w(0) = {wh/w0:.9f}")
    # weights sum to 1 in the limit
    alpha = 0.1
    s = sum(alpha * (1 - alpha) ** i for i in range(10000))
    check("P11 EMA weights sum to 1", close(s, 1.0, 1e-9), f"sum = {s:.12f}")


# ------------------------------------------------------------ P12 post-filter
def p12_postfilter_overfetch():
    """Post-filtering needs K >= k/s in expectation; verified by simulation."""
    rng = random.Random(7)
    n, k = 200000, 10
    for s in (0.5, 0.1, 0.01, 0.001):
        K = math.ceil(k / s)
        # label each of the top-K as matching with probability s
        got = [sum(1 for _ in range(K) if rng.random() < s) for _ in range(400)]
        mean = sum(got) / len(got)
        check(f"P12 K=k/s yields ~k matches at s={s}", abs(mean - k) < 1.0,
              f"K={K}, mean matches {mean:.2f} (target {k})")
        # and roughly half the time you fall short -- which is why real systems over-fetch
    frac_of_corpus = math.ceil(k / 0.0001) / 1e6
    check("P12 at s=1e-4 you scan 10% of a 1M corpus",
          abs(frac_of_corpus - 0.1) < 1e-9, f"{frac_of_corpus*100:.1f}%")


# ------------------------------------------------------------ P13 amplification
def p13_amplification():
    """Leveled: W ~ T*L, R ~ L, S ~ 1+1/T.  Size-tiered: W ~ L, R ~ T*L, S ~ 2."""
    T, base = 10, 64e6
    for data in (1e9, 8e9, 64e9, 512e9):
        L = max(1, math.ceil(math.log(data / base, T)))
        lev = dict(W=T * L + 1, R=L + 1, S=1 + 1 / T)
        tier = dict(W=L + 1, R=T * L, S=2.0)
        check(f"P13 leveled trades write for read at {data/1e9:.0f} GB",
              lev["W"] > tier["W"] and lev["R"] < tier["R"] and lev["S"] < tier["S"],
              f"L={L}  leveled W/R/S = {lev['W']:.0f}/{lev['R']:.0f}/{lev['S']:.2f}   "
              f"tiered = {tier['W']:.0f}/{tier['R']:.0f}/{tier['S']:.2f}")
    # the RUM statement: no strategy dominates on all three
    check("P13 neither strategy dominates on all three axes", True,
          "leveled wins R and S, size-tiered wins W -- that is the conjecture, felt")


# ------------------------------------------------------------ P14 tail at scale
def p14_tail_at_scale():
    """P(at least one of N slow) = 1 - (1-p)^N, verified by simulation."""
    rng = random.Random(8)
    p = 0.01
    for N in (1, 10, 100, 500):
        closed = 1 - (1 - p) ** N
        trials = 40000
        hits = sum(1 for _ in range(trials)
                   if any(rng.random() < p for _ in range(N)))
        emp = hits / trials
        check(f"P14 tail at N={N}", abs(emp - closed) < 0.02,
              f"closed {closed*100:.2f}%, simulated {emp*100:.2f}%")


# ------------------------------------------------------------ P15 systolic reuse
def p15_systolic_reuse():
    """A k x k weight-stationary array: O(k) operands fetched per k^2 MACs."""
    for k in (8, 64, 256):
        fetched = 2 * k          # one column of activations + one row of psums
        macs = k * k
        reuse = macs / fetched
        check(f"P15 operand reuse is O(k) at k={k}", close(reuse, k / 2, 1e-9),
              f"{macs} MACs per {fetched} operands = {reuse:.1f}x reuse")
    # TPUv1 headline figure from two integers
    k, clock = 256, 700e6
    tops = 2 * k * k * clock / 1e12
    check("P15 TPUv1 92 TOPS from k=256 at 700 MHz", abs(tops - 91.75) < 0.1,
          f"2 * {k}^2 * {clock/1e6:.0f}MHz = {tops:.2f} TOPS (reported: 92)")


# ------------------------------------------------- P16 concentration of distances
def p16_concentration():
    """Relative contrast -> 1 as dimension grows, for i.i.d. coordinates."""
    rng = random.Random(9)

    def rc(d, n=1500):
        pts = [[rng.gauss(0, 1) for _ in range(d)] for _ in range(n)]
        norms = [math.sqrt(sum(x * x for x in p)) for p in pts]
        pts = [[x / nm for x in p] for p, nm in zip(pts, norms)]
        q = [rng.gauss(0, 1) for _ in range(d)]
        qn = math.sqrt(sum(x * x for x in q)); q = [x / qn for x in q]
        ds = sorted(math.sqrt(max(0.0, 2 - 2 * sum(a * b for a, b in zip(p, q))))
                    for p in pts)
        return (sum(ds) / len(ds)) / ds[0]

    vals = [(d, rc(d)) for d in (2, 8, 64, 512)]
    monotone = all(vals[i][1] > vals[i + 1][1] for i in range(len(vals) - 1))
    check("P16 relative contrast decreases with dimension", monotone,
          "  ".join(f"d={d}: RC={v:.2f}" for d, v in vals))
    check("P16 RC approaches 1 in high dimension", vals[-1][1] < 1.2,
          f"d=512 gives RC={vals[-1][1]:.3f}")


# ------------------------------------------------------------ P17 cosine/L2
def p17_metric_equivalence():
    """On unit vectors: ||a-b||^2 = 2 - 2<a,b>, so the three rankings coincide."""
    rng = random.Random(10)
    d, n = 32, 400
    def norm(v):
        s = math.sqrt(sum(x * x for x in v)); return [x / s for x in v]
    pts = [norm([rng.gauss(0, 1) for _ in range(d)]) for _ in range(n)]
    q = norm([rng.gauss(0, 1) for _ in range(d)])

    dot = [sum(a * b for a, b in zip(p, q)) for p in pts]
    l2 = [math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q))) for p in pts]
    ident = max(abs(l2[i] ** 2 - (2 - 2 * dot[i])) for i in range(n))
    check("P17 ||a-b||^2 = 2 - 2<a,b> on unit vectors", ident < 1e-9,
          f"max deviation {ident:.2e}")

    by_dot = sorted(range(n), key=lambda i: -dot[i])
    by_l2 = sorted(range(n), key=lambda i: l2[i])
    check("P17 max-dot and min-L2 give identical orderings", by_dot == by_l2,
          f"top-10 identical: {by_dot[:10] == by_l2[:10]}")

    # and it FAILS without normalisation -- the silent recall bug
    raw = [[rng.gauss(0, 1) * rng.uniform(0.2, 5) for _ in range(d)] for _ in range(n)]
    rdot = [sum(a * b for a, b in zip(p, q)) for p in raw]
    rl2 = [math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q))) for p in raw]
    check("P17 the equivalence FAILS on unnormalised vectors",
          sorted(range(n), key=lambda i: -rdot[i]) != sorted(range(n), key=lambda i: rl2[i]),
          "forgetting to normalise silently changes the ranking -- P02's E6")


# ------------------------------------------------------------ P18 Zipf head mass
def p18_zipf_head():
    """Share of mass in the top f fraction under a Zipf(alpha) popularity law."""
    n = 10000
    for alpha, exp_top1 in ((0.5, 9.4), (0.8, 30.0), (1.0, 53.0), (1.2, 75.1)):
        w = [1 / r ** alpha for r in range(1, n + 1)]
        tot = sum(w)
        top1 = sum(w[: n // 100]) / tot * 100
        check(f"P18 Zipf({alpha}) top 1% mass", abs(top1 - exp_top1) < 0.6,
              f"{top1:.1f}% (expected {exp_top1}%)")


def main() -> int:
    print("Verifying every derivation in proofs.md\n")
    for fn in (p1_softmax_scale, p2_bloom_optimal_k, p4_quorum_intersection,
               p5_littles_law, p6_decode_intensity, p7_ad_modes, p8_matmul_backward,
               p9_sample_size, p10_peeking, p11_ema_halflife,
               p12_postfilter_overfetch, p13_amplification, p14_tail_at_scale,
               p15_systolic_reuse, p16_concentration, p17_metric_equivalence,
               p18_zipf_head):
        fn()
    passed = sum(1 for _, ok, _ in RESULTS if ok)
    failed = len(RESULTS) - passed
    print(f"\n{passed}/{len(RESULTS)} checks passed" +
          (f", {failed} FAILED" if failed else ""))
    if failed:
        for n, ok, d in RESULTS:
            if not ok:
                print(f"  FAILED: {n}  {d}")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())