"""Assemble initialization_provenance from recorded client events (server get_info + fixed-batch forward) and validate."""
import json, sys, os
sys.path.insert(0, os.path.expanduser("~/repos/research-exchange"))
from research_exchange import restoration
spec = json.load(open("spec.json"))
probe = json.load(open("probe_results.json"))
arms = []
for R in (1, 32):
    for LR in ("1e-4", "3e-4", "1e-3"):
        d = f"runs/r{R}_lr{LR}"
        ev = [json.loads(l) for l in open(f"{d}/client_events.jsonl")]
        init = [e for e in ev if e["kind"] == "init_load"]
        restores = [e for e in ev if e["kind"] == "resume_restore"]
        e0 = init[0]
        checks = [{"after_restore": i, "source": "server_response", "actual": e["info"]["optimizer_config"]["type"],
                   "actual_config": {k: v for k, v in e["info"]["optimizer_config"].items() if k != "type"},
                   "server_lora_rank": e["info"]["lora_rank"], "before_next_update": True,
                   "evidence_ref": f"{d}/client_events.jsonl#{e['kind']}"} for i, e in enumerate([e0] + restores)]
        arms.append({"name": f"rank{R}_lr{LR}", "restore_count": len(restores),
                     "initialization": {"base_checkpoint": "Qwen/Qwen2.5-3B", "source": "server_response",
                                        "checkpoint_uri": e0["requested_path"], "checkpoint_immutable": True,
                                        "checkpoint_scope": "full_initialization",
                                        "load_note": "Load acknowledged only as a client created 'via load_weights'; the server response does not echo the loaded URI, so loaded_checkpoint_uri is not recorded. save_weights_external (for an adapter hash) wedged the client in the probe, so adapter_sha256 is not measured.",
                                        "evidence_ref": f"{d}/client_events.jsonl", "before_first_update": True},
                     "optimizer_checks": checks,
                     "initial_output": {"available": True, "batch_sha256": e0["batch_sha256"], "output_sha256": e0["output_sha256"],
                                        "evidence_ref": f"{d}/client_events.jsonl", "before_first_update": True}})
prov = {"version": 1, "arms": arms,
        "probe": {"evidence_ref": "probe_results.json", "fresh_rank1_vs_rank32_identical": probe["fresh_outputs_identical_across_ranks"],
                  "quiet_restore_diagnostic": "diag_init.json (4 fresh restores, all output_sha256 5dd1397c..., max abs diff 0.0)"}}
val = restoration.validate(spec, prov)
json.dump({"initialization_provenance": prov, "validation": val}, open("restoration_validation.json", "w"), indent=1)
print(val["status"], val["reason"])
for c in val["checks"]:
    if c["status"] != "passed": print(" ", c["status"], c["name"] if "name" in c else "", c["reason"][:160])