"""
roofline.py: measure this computer's two speed limits, then put real
operations on the roofline and see how close each one gets to its roof.
Plain NumPy, float32. Every test runs several times; we keep the median.
"""
import time
import numpy as np

def timed(fn, reps=5):
    fn()                                              # warm-up
    ts = []
    for _ in range(reps):
        t = time.perf_counter(); fn(); ts.append(time.perf_counter() - t)
    return float(np.median(ts))

# ---- 1) The two roofs -------------------------------------------------
N = 10_000
W = np.random.rand(N, N).astype(np.float32)           # a 400 MB table
x1 = np.random.rand(N, 1).astype(np.float32)
sec = timed(lambda: W @ x1)
BW = W.nbytes / sec                                   # bytes per second

a = np.random.rand(2048, 2048).astype(np.float32)
b = np.random.rand(2048, 2048).astype(np.float32)
sec = timed(lambda: a @ b)
PEAK = 2 * 2048**3 / sec                              # FLOPs per second

RIDGE = PEAK / BW
print(f"memory roof : {BW/1e9:6.1f} GB/s")
print(f"math roof   : {PEAK/1e9:6.1f} GFLOP/s")
print(f"ridge point : {RIDGE:6.1f} FLOPs per byte\n")

def roof(intensity):                                  # the whole model, one line
    return min(PEAK, BW * intensity)

rows = []
def report(name, flops, nbytes, sec):
    I = flops / nbytes
    got = flops / sec
    rows.append((name, I, got))
    print(f"{name:<28} {I:8.2f} {got/1e9:9.1f} {roof(I)/1e9:9.1f} {100*got/roof(I):6.0f}%")

print(f"{'operation':<28} {'FLOP/B':>8} {'GFLOP/s':>9} {'roof':>9} {'of roof':>7}")

# ---- 2) Big table x B tokens (the model's weight tables) ---------------
for B in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]:
    X = np.random.rand(N, B).astype(np.float32)
    flops = 2 * N * N * B
    nbytes = 4 * (N * N + N * B + N * B)              # read table + tokens, write result
    sec = timed(lambda: W @ X, reps=5 if B < 128 else 3)
    report(f"table x {B} tokens", flops, nbytes, sec)

# ---- 3) Attention while typing: each user reads their OWN notebook ----
T, D = 8_192, 768                                     # 8,192 notes, 25 MB per user
for users in [16, 64]:
    K = [np.random.rand(T, D).astype(np.float32) for _ in range(users)]
    q = [np.random.rand(D).astype(np.float32) for _ in range(users)]
    flops = 2 * T * D * users
    nbytes = 4 * T * D * users
    sec = timed(lambda: [k @ v for k, v in zip(K, q)])
    report(f"notebook read, {users} users", flops, nbytes, sec)
    del K

# ---- 4) Small steps: add two long lists of numbers --------------------
u = np.random.rand(50_000_000).astype(np.float32)
v = np.random.rand(50_000_000).astype(np.float32)
out = np.empty_like(u)
sec = timed(lambda: np.add(u, v, out=out))
report("add two lists", u.size, 12 * u.size, sec)

np.save("roofline_points.npy", np.array([(r[1], r[2]) for r in rows]))
np.save("roofline_roofs.npy", np.array([BW, PEAK]))