"""Throughput vs probe_chunk for ZBP steps (informs the H200 config)."""
import os, sys, time, argparse
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
import torch, torch.nn as nn
from zbp_scaling.model import ScalingLM
from zbp_scaling.zbp import ZBPConfig, zbp_blocks
from zbp_scaling.zbp.autograd import seed_probes

p = argparse.ArgumentParser()
p.add_argument("--device", default="cuda:3")
p.add_argument("--d", type=int, default=512); p.add_argument("--layers", type=int, default=8)
p.add_argument("--vocab", type=int, default=8192); p.add_argument("--seq", type=int, default=1024)
p.add_argument("--bs", type=int, default=8)
a = p.parse_args()
dev = torch.device(a.device); seed_probes(0, dev)
x = torch.randint(0, a.vocab, (a.bs, a.seq), device=dev); y = torch.randint(0, a.vocab, (a.bs, a.seq), device=dev)

def bench(mode, n, chunk, iters=4):
    cfg = ZBPConfig(mode=mode, n_probes=n, eps=0.1, probe_chunk=chunk)
    m = ScalingLM(a.vocab, a.d, a.layers, 8, a.seq, cfg=cfg).to(dev)
    opt = torch.optim.AdamW(m.parameters(), lr=1e-4)
    for _ in range(2):
        loss = nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()
    torch.cuda.synchronize(dev); t = time.time()
    for _ in range(iters):
        loss = nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()
    torch.cuda.synchronize(dev)
    dt = (time.time() - t) / iters
    print(f"{mode:3s} n={n:3d} chunk={chunk:3d}: {dt*1000:7.0f} ms/step  {a.bs*a.seq/dt/1000:7.1f} ktok/s  peakmem {torch.cuda.max_memory_allocated(dev)/2**30:.1f} GB")
    torch.cuda.reset_peak_memory_stats(dev)
    return dt

t_bp = bench("bp", 0, 0)
for n in (16, 64):
    for chunk in (8, 16, 32, 64):
        if chunk > 2 * n: continue
        dt = bench("cd", n, chunk)
        print(f"    -> multiplier vs BP: {dt/t_bp:.1f}x")