import json, torch, math
from analyze_preflight import load, diff, blocks, split, isB, err, sha
from dimuon_candidate import candidate_step
def allclose(actual, ref, atol, rtol):
    a = torch.cat([actual[k].double().flatten() for k in sorted(ref)]); r = torch.cat([ref[k].double().flatten() for k in sorted(ref)])
    ok = (a - r).abs() <= atol + rtol * r.abs()
    return {"elementwise_pass_fraction": ok.double().mean().item(), "n_elements": ok.numel(), "n_fail": int((~ok).sum())}
res = json.load(open("preflight_results.json"))
P0 = load("pf-init"); out = {}
# AdamW: identical grads/state from the AdamW client's own saved exp_avg; reference = stock torch.optim.AdamW
A1 = load("pf-adam-1e-4"); mA = load("pf-adam-1e-4", "optimizer.state.exp_avg."); vA = load("pf-adam-1e-4", "optimizer.state.exp_avg_sq.")
ref = {}
for k in P0:
    p = torch.nn.Parameter(P0[k].clone()); p.grad = mA[k] / 0.1
    o = torch.optim.AdamW([p], lr=1e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0); o.step()
    ref[k] = p.detach() - P0[k]
dA = diff(P0, A1)
out["baseline_update"] = {**err(dA, ref), **allclose(dA, ref, 1e-6, 1e-3), "reference": f"torch.optim.AdamW (torch {torch.__version__}) on the fork client's own step-1 exp_avg/(1-beta1) gradient and identical fp32 init"}
out["adam_v_consistency_max_rel"] = max(float(((vA[k] / 0.05).sqrt() - (mA[k] / 0.1).abs()).abs().max() / max(float((mA[k] / 0.1).abs().max()), 1e-30)) for k in P0)
# Cross-client gradient noise floor: G (eps-trick) vs A's exact gradient
uG = {k: -v for k, v in diff(P0, load("pf-grad")).items()}; gG = {k: u * 1e6 / (1e6 - u.abs()) for k, u in uG.items()}
out["cross_client_grad_agreement"] = err(gG, {k: mA[k] / 0.1 for k in P0})
# Dimuon: identical grads from the Dimuon client's own BF16 step-1 momentum (= bf16((1-beta1) g))
D1 = load("pf-dim-0.001"); mD = load("pf-dim-0.001", "optimizer.state.momentum_buffer.")
def ref_dimuon(gd, lr):
    r = {}
    for k, v in gd.items():
        if not isB(k): r[k] = torch.zeros_like(v); continue
        outs = [candidate_step(torch.zeros_like(gb), gb, None, learning_rate=lr, beta1=0.9, damping=1e-4, rank_lr_correction_exponent=0.175)[0] for gb in blocks(k, v)]
        gg, s = split(k, v.shape[0]); r[k] = outs[0] if s is None else torch.cat([o.reshape(gg, -1, v.shape[1]) for o in outs], 1).reshape(v.shape)
    return r
dD = diff(P0, D1)
for tag, gd in (("own_bf16_momentum", {k: mD[k].float() / 0.1 for k in P0}), ("other_client_fp32_grad", {k: mA[k] / 0.1 for k in P0})):
    r = ref_dimuon(gd, 1e-3); out[f"reference_update_{tag}"] = {**err(dD, r), **allclose(dD, r, 1e-5, 1e-2)}
# Dimuon same-state cross-client noise floor: 10x lr1e-4 update vs lr1e-3 update (separate clients, identical P0, batch)
d4 = {k: 10 * v for k, v in diff(P0, load("pf-dim-0.0001")).items()}; d2 = {k: v / 10 for k, v in diff(P0, load("pf-dim-0.01")).items()}
out["lr_scaled_update_1e-4_x10_vs_1e-3"] = {**err(d4, dD), **allclose(d4, dD, 1e-5, 1e-2)}
out["lr_scaled_update_1e-2_div10_vs_1e-3"] = {**err(d2, dD), **allclose(d2, dD, 1e-5, 1e-2)}
U2, R2 = load("pf-U2"), load("pf-R2"); dU, dR = diff(D1, U2), diff(D1, R2)
out["restore_replay_update"] = {**err(dR, dU), **allclose(dR, dU, 1e-5, 1e-2)}
mU2 = load("pf-U2", "optimizer.state.momentum_buffer."); mR2 = load("pf-R2", "optimizer.state.momentum_buffer.")
out["restore_replay_momentum"] = err(mR2, mU2)
res["v2"] = out; json.dump(res, open("preflight_results.json", "w"), indent=1); print(json.dumps(out, indent=1))