# Identical-input parity: deployed Miles Dimuon code (packed fused B layouts) vs Aaron's reference (unpacked per-projection blocks).
import sys, json, types, importlib.util, torch, hashlib
from analyze_preflight import load, diff, blocks, split, isB, err
from analyze2 import allclose
from dimuon_candidate import candidate_step
src = "/home/ubuntu/repos/miles/miles/backends/megatron_utils/lora/dimuon.py"
for name in ("miles", "miles.backends", "miles.backends.megatron_utils", "miles.backends.megatron_utils.lora"):
    sys.modules.setdefault(name, types.ModuleType(name))
g = importlib.util.spec_from_file_location("miles.backends.megatron_utils.lora.gdn_utils", src.replace("dimuon.py", "gdn_utils.py"))
gm = importlib.util.module_from_spec(g); sys.modules[g.name] = gm; g.loader.exec_module(gm)
s = importlib.util.spec_from_file_location("miles.backends.megatron_utils.lora.dimuon", src)
dm = importlib.util.module_from_spec(s); sys.modules[s.name] = dm; s.loader.exec_module(dm)
P0 = load("pf-init"); g1 = {k: v / 0.1 for k, v in load("pf-adam-1e-4", "optimizer.state.exp_avg.").items()}
uG = {k: -v for k, v in diff(P0, load("pf-grad")).items()}; g2 = {k: u * 1e6 / (1e6 - u.abs()) for k, u in uG.items()}
Bk = sorted(k for k in P0 if isB(k))
def miles_run(lr, grads):
    ps = {k: torch.nn.Parameter(P0[k].clone()) for k in Bk}
    opt = dm.Dimuon(list(ps.values()), lr=lr, momentum=0.9)
    opt.set_specs({id(ps[k]): dm.FactorSpec(rank=32, row_split=split(k, ps[k].shape[0])[1] if "linear_qkv" not in k else (8 * (ps[k].shape[0] // 20), ps[k].shape[0] // 20, ps[k].shape[0] // 20)) for k in Bk})
    outs = []
    for gr in grads:
        before = {k: ps[k].detach().clone() for k in Bk}
        for k in Bk: ps[k].grad = gr[k].clone()
        opt.step(); outs.append({k: ps[k].detach() - before[k] for k in Bk})
    return outs
def ref_run(lr, grads):
    par = {k: [b.clone() for b in blocks(k, P0[k])] for k in Bk}; mom = {k: None for k in Bk}; outs = []
    for gr in grads:
        upd = {}
        for k in Bk:
            gb = blocks(k, gr[k]); new, nm = [], []
            for i, (pb, gbb) in enumerate(zip(par[k], gb)):
                n, m = candidate_step(pb, gbb, None if mom[k] is None else mom[k][i], learning_rate=lr, beta1=0.9, damping=1e-4, rank_lr_correction_exponent=0.175)
                new.append(n); nm.append(m)
            d = [n - p for n, p in zip(new, par[k])]; par[k], mom[k] = new, nm
            gg, sp = split(k, P0[k].shape[0])
            upd[k] = d[0] if sp is None else torch.cat([x.reshape(gg, -1, P0[k].shape[1]) for x in d], 1).reshape(P0[k].shape)
        outs.append(upd)
    return outs
res = {"miles_source_sha256": hashlib.sha256(open(src, "rb").read()).hexdigest(), "reference_sha256": hashlib.sha256(open("dimuon_candidate.py", "rb").read()).hexdigest()}
M = miles_run(1e-3, [g1, g2]); R = ref_run(1e-3, [g1, g2])
res["reference_update_step1_fresh"] = {**err(M[0], R[0]), **allclose(M[0], R[0], 1e-5, 1e-2)}
res["reference_update_step2_stale_scales"] = {**err(M[1], R[1]), **allclose(M[1], R[1], 1e-5, 1e-2)}
# Deployed server update vs local Miles on its own BF16 momentum-derived gradient (server fidelity, lossy input)
mD = load("pf-dim-0.001", "optimizer.state.momentum_buffer."); dD = diff(P0, load("pf-dim-0.001"))
Ms = miles_run(1e-3, [{k: mD[k].float() / 0.1 for k in Bk}])[0]
res["server_vs_local_miles_bf16_momentum_input"] = err({k: dD[k] for k in Bk}, Ms)
res["lr_scaling_local"] = [{"learning_rate": lr, "update_norm": float(torch.cat([v.double().flatten() for v in miles_run(lr, [g1])[0].values()]).norm())} for lr in (1e-4, 1e-3, 1e-2)]
json.dump(res, open("local_parity.json", "w"), indent=1); print(json.dumps(res, indent=1))