import glob, json, hashlib, math, sys, torch, numpy as np
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
from dimuon_candidate import candidate_step
torch.set_float32_matmul_precision("highest")
PFX = "optimizer.state.fp32_param."
def load(name, prefix=PFX):
    p = glob.glob(f"ckpt/{name}/*/miles")[0]
    md = dcp.FileSystemReader(p).read_metadata()
    sd = {k: torch.empty(v.size, dtype=v.properties.dtype) for k, v in md.state_dict_metadata.items() if k.startswith(prefix) and hasattr(v, "size")}
    dcp.load(sd, storage_reader=dcp.FileSystemReader(p), planner=DefaultLoadPlanner(allow_partial_load=True))
    return {k[len(prefix):]: v for k, v in sd.items()}
def sha(sd):
    h = hashlib.sha256()
    for k in sorted(sd): h.update(k.encode()); h.update(sd[k].contiguous().view(torch.uint8).numpy().tobytes())
    return h.hexdigest()
def split(name, rows):
    if "linear_qkv" in name:
        hd = rows // 20; assert hd * 20 == rows; return 2, (8 * hd, hd, hd)
    if "linear_fc1" in name: return 1, (rows // 2, rows // 2)
    return 1, None
def blocks(name, t):
    g, s = split(name, t.shape[0])
    if s is None: return [t]
    return [b.reshape(-1, t.shape[1]) for b in torch.split(t.reshape(g, sum(s), t.shape[1]), s, dim=1)]
def err(actual, ref):
    a = torch.cat([actual[k].double().flatten() for k in sorted(ref)]); r = torch.cat([ref[k].double().flatten() for k in sorted(ref)])
    rn = r.norm().item(); an = a.norm().item()
    return {"max_abs_error": (a - r).abs().max().item(), "reference_max_abs": r.abs().max().item(),
            "relative_l2_error": (a - r).norm().item() / rn if rn else None,
            "cosine_similarity": (a @ r).item() / (an * rn) if an and rn else None}
def diff(a, b): return {k: b[k] - a[k] for k in a}
P0 = load("pf-init"); res = {"init_fp32_sha256": sha(P0)}
res["init_bf16_adapter_sha256"] = sha(load("pf-init", prefix="language_model."))
isB = lambda k: k.endswith("linear_out.weight")
res["init_B_all_zero"] = all(float(P0[k].abs().max()) == 0 for k in P0 if isB(k))
uG = {k: -v for k, v in diff(P0, load("pf-grad")).items()}  # u = lr*g/(|g|+eps), lr=eps=1e6
EPS = 1e6; g = {k: u * EPS / (EPS - u.abs()) for k, u in uG.items()}
res["recovered_grad_norm"] = math.sqrt(sum(float((v.double() ** 2).sum()) for v in g.values()))
res["recovered_grad_A_max_abs"] = max(float(v.abs().max()) for k, v in g.items() if not isB(k))
dA = diff(P0, load("pf-adam-1e-4"))
refA = {k: -1e-4 * v / (v.abs() + 1e-8) for k, v in g.items()}
res["baseline_update"] = err(dA, refA)
dD = {lr: diff(P0, load(f"pf-dim-{lr:g}")) for lr in (1e-4, 1e-3, 1e-2)}
refD = {}
for k, v in g.items():
    if not isB(k): refD[k] = torch.zeros_like(v); continue
    outs = []
    for gb in blocks(k, v):
        new, _ = candidate_step(torch.zeros_like(gb), gb, None, learning_rate=1e-3, beta1=0.9, damping=1e-4, rank_lr_correction_exponent=0.175)
        outs.append(new)
    gg, s = split(k, v.shape[0])
    refD[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)
res["reference_update"] = err(dD[1e-3], refD)
res["reference_update_B_only"] = err({k: dD[1e-3][k] for k in refD if isB(k)}, {k: refD[k] for k in refD if isB(k)})
res["dimuon_A_step1_max_abs"] = max(float(dD[1e-3][k].abs().max()) for k in refD if not isB(k))
res["lr_scaling"] = [{"learning_rate": lr, "update_norm": math.sqrt(sum(float((v.double() ** 2).sum()) for v in dD[lr].values()))} for lr in (1e-4, 1e-3, 1e-2)]
r0 = res["lr_scaling"][0]["update_norm"] / 1e-4
res["lr_scaling_max_rel_change"] = max(abs(x["update_norm"] / x["learning_rate"] / r0 - 1) for x in res["lr_scaling"])
def srank(d):
    vals = []
    for k, v in d.items():
        if not isB(k): continue
        for b in blocks(k, v):
            s = torch.linalg.svdvals(b.double()); vals.append(float((s ** 2).sum() / (s[0] ** 2 * 32)))
    return float(np.mean(vals)), len(vals)
res["dimuon_B_stable_rank_step1"], res["n_B_blocks"] = srank(dD[1e-3])
res["adamw_B_stable_rank_step1"], _ = srank(dA)
U2, R2 = load("pf-U2"), load("pf-R2")
res["restore_replay_params"] = err(diff(load("pf-dim-0.001"), R2), diff(load("pf-dim-0.001"), U2))
json.dump(res, open("preflight_results.json", "w"), indent=1); print(json.dumps(res, indent=1))