import json, hashlib, glob, os, time
rd = lambda f: [json.loads(l) for l in open(f)]
P, S, C = rd("pilot_log.jsonl"), rd("select_log.jsonl"), rd("confirm_log.jsonl")
F = json.load(open("final_numbers.json")); res = F["res"]; sel = json.load(open("selection_frozen.json"))
RC = json.load(open("restore_check.json")); CP = json.load(open("correctness_preflight.json"))
spec = json.load(open("spec.json")); T_START = 1793741580  # session start 21:33 UTC (deploy began 1793741987)
cleanup = json.load(open("cleanup.json")) if os.path.exists("cleanup.json") else None
def st(L, s): return [e for e in L if e["stage"] == s]
# batch-hash pairing in confirmation
steps = {}
for e in st(C, "train_step"): steps.setdefault((e["arm"].split("_s")[-1], e["step"]), set()).add(e["batch_sha256"])
pair_match = sum(len(v) == 1 for v in steps.values()) / len(steps)
pt = {e["params_type"] for e in st(C, "train_step") + st(S, "train_step")}
io = {e["arm"]: e for e in st(C, "initial_output") + st(S, "initial_output")}
gi = {}
for e in st(C, "get_info") + st(S, "get_info"): gi.setdefault(e["arm"], []).append(e)
s0 = {e["arm"]: e["path"] for e in st(C, "step0_saved")}
arms_prov = []
for a in json.load(open("select_arms.json")) + json.load(open("confirm_arms.json")):
    n = a["arm"]; ph = "confirm" if n in s0 else "select"
    init_name = f"{ph}-init-{a['init']}"
    r = {"arm": n, "phase": ph, "base_checkpoint": "Qwen/Qwen3.5-2B", "initialization_checkpoint": io[n]["init_path"],
         "before_first_update": True, "restore_count": io[n]["restores"],
         "optimizer_checks": [{"restores": g["restores"], "optimizer": g["optimizer"], "model": g["model"], "rank": g["rank"]} for g in gi[n]],
         "initial_output_sha256": io[n]["output_sha256"], "evidence_ref": f"{ph}_log.jsonl"}
    if n in s0:
        k = f"confirm-{n}-step0"; r.update(step0_checkpoint=s0[n], adapter_sha256=RC[k]["adapter_bf16_sha256"], fp32_param_sha256=RC[k]["fp32_param_sha256"],
              init_adapter_sha256=RC[init_name]["adapter_bf16_sha256"], adapter_matches_init=RC[k]["adapter_bf16_sha256"] == RC[init_name]["adapter_bf16_sha256"],
              server_checkpoint_optimizer=RC[k]["metadata"]["optimizer"], evidence_ref="restore_check.json")
    else:
        r.update(adapter_sha256=None, adapter_note="selection arms all restored from one saved state select-init-s0; adapter hash not downloaded per arm, only initial output hash")
    arms_prov.append(r)
pairs_ok = all(io[f"baseline_adamw_s{s}"]["output_sha256"] == io[f"treatment_dimuon_s{s}"]["output_sha256"] and RC[f"confirm-baseline_adamw_s{s}-step0"]["adapter_bf16_sha256"] == RC[f"confirm-treatment_dimuon_s{s}-step0"]["adapter_bf16_sha256"] for s in (0, 1, 2))
ph = lambda L: (st(L, "setup_done")[0]["seconds"], st(L, "train_done")[0]["seconds"], st(L, "eval_done")[0]["seconds"])
pp, sp, cp = ph(P), ph(S), ph(C)
now = cleanup["t"] if cleanup else time.time(); wall = now - T_START; gpu_h = (now - 1793742009) / 3600  # 1 trainer H100 from deploy until cleanup
src = "fork_config.py trainer_gpus_per_node=1 (single trainer GPU; inference replicas 0)"
meas = lambda sec, u, pu, smp: {"seconds": sec, "units": u, "planned_units": pu, "samples": smp, "gpu_count": 1, "allocation_source": src}
feas = {"config": spec["feasibility"]["effective_config"],
        "effective_config": {**spec["feasibility"]["effective_config"], "n_ifeval_secondary": 0, "pilot_arms": 2, "pilot_steps": 10,
                             "note": "IFEval dropped (downsizing step 1); LR grid, selection steps, seeds, test NLL unchanged"},
        "primary_endpoint": "no_robots_test_token_mean_nll_at_step_100", "source": "pilot_log.jsonl (10-step 2-arm representative pilot); select_log.jsonl and confirm_log.jsonl measured",
        "measurements": {"setup_pilot": meas(pp[0], 2, 2, 2), "training_pilot": meas(pp[1], 20, 20, 640), "heldout_eval_pilot": meas(pp[2], 1000, 1000, 1000),
                         "setup_selection": meas(sp[0], 6, 6, 6), "training_selection": meas(sp[1], 240, 240, 7680), "heldout_eval_selection": meas(sp[2], 3000, 3000, 3000),
                         "setup_confirmation": meas(cp[0], 6, 6, 6), "training_confirmation": meas(cp[1], 600, 600, 19200), "heldout_eval_confirmation": meas(cp[2], 3000, 3000, 3000)},
        "generation": None, "recovery": {"seconds": max(0.0, 150 * 60 - wall), "gpu_count": 1, "allocation_source": src,
                                         "note": f"wall used {wall/60:.1f} of 150 min; trainer GPU-h used ~{gpu_h:.2f} of 2.0 (deploy to now, 1 GPU)"}}
eid = F["eid"]
loss_eval = {"aggregation": "token_mean", "arms": F["loss_eval"]}
metrics = {"arms": F["arms"]}
if "mean" in res:
    metrics.update(primary_mean_adamw_minus_dimuon=res["mean"], primary_ci95_low=res["lo"], primary_ci95_high=res["hi"], primary_sd=res["sd"], margin=0.02)
for i, d in enumerate(res["diffs_adamw_minus_dimuon"]): metrics[f"diff_adamw_minus_dimuon_s{i}"] = d
metrics.update(selected_lr_adamw=sel["adam"]["selected_lr"], selected_lr_dimuon=sel["dimuon"]["selected_lr"], paired_batch_hash_match_confirmation=pair_match,
               wall_minutes=wall / 60, trainer_gpu_hours=gpu_h)
files = []
def add(name, kind, path=None):
    b = open(path or name, "rb").read(); files.append({"name": name, "kind": kind, "content": b.decode(), "sha256": hashlib.sha256(b).hexdigest()})
for n in ("run.py", "common.py", "fork_config.py", "preflight.py", "analyze_preflight.py", "analyze2.py", "local_parity.py", "build_preflight.py", "launch_confirm.py", "final_analysis.py", "restore_check.py", "build_output.py", "smoke.py", "probe_shared.py"): add(n, "script")
add("launch-instructions.md", "instructions", "../launch-instructions.md")
json.dump(feas["effective_config"], open("effective-config.json", "w"), indent=1); add("effective-config.json", "effective_config")
open("code-revision.txt", "w").write("fork app: spindle-fork-run_739abc9a8b (env research-exchange-runs)\nmodal-projects/spindle 5148cb62a019862ef9b50d76b7be55cddb288e8a (branch claude/dimuon-spindle-wiring, PR #85) + spindle_fork.patch\nmodal-projects/miles 228b59d90074e6cec48090e033202c5a4aff3ff3 (branch dimuon-lite, PR #17), mounted via SPINDLE_MILES_TREE; server checkpoint metadata reports miles revision fdd808c008175d4895fd8c57c95ef4fa5bb6a2a0\nshared Spindle (not used for arms): 0.1.5 @ 4af4dd2c46c009fcf8661e625195d4bf6c5dde81, rejects dimuon\nAaron reference: dimuon_candidate.py sha256 " + hashlib.sha256(open("dimuon_candidate.py", "rb").read()).hexdigest() + "\n")
add("code-revision.txt", "code_revision"); add("spindle_fork.patch", "code_patch"); add("dimuon_candidate.py", "script")
add("predictions-confirm-test.jsonl", "predictions"); add("scored-observations-final.jsonl", "metrics")
for n in ("select_log.jsonl", "confirm_log.jsonl", "pilot_log.jsonl", "correctness_preflight.json", "preflight_results.json", "local_parity.json", "restore_check.json", "selection_frozen.json", "final_numbers.json"): add(n, "metrics")
if cleanup: add("cleanup.json", "metrics")
ck = [{"arm": e["arm"], "step": e["step"], "uri": e["final_state"]} for e in st(C, "eval") + st(S, "eval")]
edge = sel["adam"]["edge"] or sel["dimuon"]["edge"]
complete = "mean" in res
verdict = "inconclusive" if (edge or not complete or not (res["lo"] > 0.02 or res["hi"] < 0.02)) else ("supported" if res["lo"] > 0.02 else "refuted")
out = {"status": "succeeded" if complete else "failed",
       "summary": (f"On fork spindle-fork-run_739abc9a8b, 100-step rank-32 LoRA SFT of Qwen3.5-2B on no_robots gave AdamW-minus-Dimuon test token-mean NLL {res['mean']:+.4f} nats (95% paired-t CI [{res['lo']:+.4f}, {res['hi']:+.4f}], 3 seeds) at selected LRs AdamW {sel['adam']['selected_lr']:g} / Dimuon {sel['dimuon']['selected_lr']:g}. "
                   f"Verdict inconclusive because both selected LRs sit at the upper grid edge (spec rule); the CI itself lies " + ("entirely below" if res["hi"] < 0.02 else "entirely above" if res["lo"] > 0.02 else "across") + " the 0.02-nat margin.") if complete else "Confirmation incomplete.",
       "metrics": metrics, "completed_arms": [e["arm"] for e in st(S, "eval") + st(C, "eval")],
       "verdict": verdict, "outcome_reason": "insufficient_evidence" if verdict == "inconclusive" else verdict,
       "deviations": [{"field": "evaluator", "original": "Secondary IFEval strict accuracy on 541 prompts (descriptive only)", "effective": "IFEval dropped; primary 500-item no_robots test token-mean NLL unchanged",
                       "reason": "Spec downsizing order step 1: measured 6-arm step times (~7-48 s) left insufficient GPU-hour/wall reserve to also run inference sampling for 6 adapters x 541 prompts."}],
       "original_hypothesis_tested": True, "correctness_preflight": CP,
       "activation_calibration": {"note": "Fork-measured before full training: stable rank from preflight step-1 Dimuon LoRA-B update (lr 1e-3, seed-0 batch 1, downloaded checkpoints, mean over 97 B blocks); NLL drop = step1 - step10 train token-mean NLL of 10-step Dimuon lr 1e-3 pilot (2.360->1.969); batch hash match over pilot pairs; step types logged client-side.",
                                  "values": {"executed_base_model": "Qwen/Qwen3.5-2B", "executed_optimizer_step_types": ",".join(sorted(pt)), "dimuon_B_update_normalized_stable_rank_step1": 0.7930079444363333,
                                             "batch_hash_match": 1.0, "dimuon_10step_pilot_train_nll_drop": 0.391, "confirmation_paired_batch_hash_match": pair_match}},
       "initialization_provenance": {"version": 1, "arms": arms_prov, "paired_init_matched_all_seeds": pairs_ok},
       "feasibility": feas, "loss_eval": loss_eval,
       "evaluation": {"version": 1, "evaluators": {eid: F["prov"]}, "scores": F["scores"], "comparisons": [{"metric": "no_robots_test_token_mean_nll_at_step_100", "arms": ["baseline", "treatment"]}]},
       "decision": {"version": 1, "observations_file": {"name": "scored-observations-final.jsonl", "sha256": F["obs_sha"]},
                    "selection": {"frozen_before_confirmation": True, "selection_split": "500 held-out no_robots train rows (Random(1234) shuffle, first 500), disjoint from training rows", "selection_steps": 40,
                                  "adamw": sel["adam"], "dimuon": sel["dimuon"], "grid_edge_selected": edge}},
       "evidence": {"version": 1, "files": files, "checkpoints": ck, "per_arm_metrics": F["arms"]},
       "adaptation_manifests": []}
if cleanup: out["cleanup"] = cleanup
json.dump(out, open("final_output.json", "w"), allow_nan=False)
print(verdict, pairs_ok, pair_match, sum(len(f["content"]) for f in files) / 1e6, "MB")