import json, hashlib, torch
from loadlib import load
spec = json.load(open("spec.json")); sp = spec["correctness_preflight"]["paths"][0]
pr, lp = json.load(open("preflight_results.json")), json.load(open("local_parity.json"))
g1 = load("pf-adam-1e-4", "optimizer.state.exp_avg.")
h = hashlib.sha256()
for k in sorted(g1): h.update(k.encode()); h.update(g1[k].contiguous().view(torch.uint8).numpy().tobytes())
log = json.load(open("preflight_log.json")); bh = [e for e in log if e["stage"] == "fixed_batch"][0]
keep = ("max_abs_error", "reference_max_abs", "relative_l2_error", "cosine_similarity")
pf = {"required": True, "paths": [{
 "id": sp["id"],
 "fixed_input": {"sha256": h.hexdigest(), "description": f"Step-1 fixture: seed-0 first no_robots train batch (32 seqs, sha256 {bh['batch1_sha256']}, cross_entropy, cookbook weights), fork init state tinker://43f94846bfb916692227a10acb4ae578/weights/pf-init (fp32 sha256 {pr['init_fp32_sha256']}, LoRA B=0), optimizer state reset per client. Hash covers the exact fp32 gradient fixture (AdamW client exp_avg). Stale-scale step 2 uses a second fixed fp32 gradient. Retrieved via modal volume get run_739abc9a8b-checkpoints."},
 "baseline_revision": sp["baseline_revision"], "reference_revision": sp["reference_revision"],
 "resolved_revisions": {
   "baseline_executed": "fork spindle-fork-run_739abc9a8b: modal-projects/spindle 5148cb62a019862ef9b50d76b7be55cddb288e8a (PR #85 branch claude/dimuon-spindle-wiring) + spindle_fork.patch (requires_proxy_auth=True); shared Spindle deployment 4af4dd2c46c009fcf8661e625195d4bf6c5dde81 is pre-#85 (rejects dimuon), so no post-#85 shared deployment exists. AdamW reference = stock torch.optim.AdamW on identical exp_avg-derived gradient and init.",
   "reference_executed": f"Aaron's DiMuon reconstruction has no git commit; pinned by content sha256 {lp['reference_sha256']} (dimuon_candidate.py, code block of Notion page 3fe1e7f169498132903de1b12bfe4205). Miles PR #17 = modal-projects/miles 228b59d90074e6cec48090e033202c5a4aff3ff3 (dimuon.py sha256 {lp['miles_source_sha256']}), mounted via SPINDLE_MILES_TREE."},
 "implementation_revision": f"spindle 5148cb62a019862ef9b50d76b7be55cddb288e8a+proxy-auth patch; miles 228b59d90074e6cec48090e033202c5a4aff3ff3 (dimuon.py sha256 {lp['miles_source_sha256']})",
 "precision": "fp32 master params and gradients; Dimuon working momentum fp32, persisted momentum bf16, Polar Express 5 iters horner_fp16; AdamW fp32 exp_avg/exp_avg_sq; model forward/backward bf16",
 "normalization": "cross_entropy sum of weighted token NLL (cookbook weights, counted by weight>0); Dimuon per-projection-block direction scaled to Frobenius sqrt(rank=32), update lr*(32/32)^0.175, damping 1e-4, beta1 0.9, no grad clip; AdamW bias-corrected first step lr*m_hat/(sqrt(v_hat)+1e-8)",
 "errors": {
   "baseline_update": {k: pr["v2"]["baseline_update"][k] for k in keep},
   "reference_update": {k: lp["reference_update_step1_fresh"][k] for k in keep},
 },
 "diagnostic_errors": {
   "reference_update_stale_scales_step2_packed_vs_unpacked": {k: lp["reference_update_step2_stale_scales"][k] for k in keep},
   "server_dimuon_vs_reference_bf16_momentum_reconstructed_grad": {k: pr["v2"]["reference_update_own_bf16_momentum"][k] for k in keep},
   "server_dimuon_vs_local_miles_same_bf16_input": lp["server_vs_local_miles_bf16_momentum_input"],
   "server_dimuon_vs_reference_other_client_fp32_grad": {k: pr["v2"]["reference_update_other_client_fp32_grad"][k] for k in keep},
   "server_cross_client_noise_floor_lr1e-4x10_vs_lr1e-3": {k: pr["v2"]["lr_scaled_update_1e-4_x10_vs_1e-3"][k] for k in keep},
   "server_cross_client_gradient_agreement": pr["v2"]["cross_client_grad_agreement"],
   "restore_replay_update_vs_uninterrupted": {k: pr["v2"]["restore_replay_update"][k] for k in keep},
 },
 "reference_update_note": "reference_update compares the deployed Miles Dimuon code (packed fused qkv/fc1 B) against Aaron's reference (unpacked per-projection blocks) on bitwise-identical fp32 inputs. The server-executed update cannot be fed an identical fp32 gradient; vs reference it shows max_abs 1.62e-5 (> 1.43e-5 threshold) with bf16-reconstructed input, identical to its error vs local Miles code on the same input, and 1.33e-5 (passes) with another client's fp32 gradient; server cross-client noise floor is 1.11e-5.",
 "lr_scaling": pr["lr_scaling"],
}]}
json.dump(pf, open("correctness_preflight.json", "w"), indent=1)