# Builds final structured output from select/confirm logs. Usage: python final_analysis.py
import json, hashlib, math, statistics, glob, os, sys
W = os.path.dirname(os.path.abspath(__file__)); os.chdir(W)
rd = lambda f: [json.loads(l) for l in open(f)]
S, C = rd("select_log.jsonl"), rd("confirm_log.jsonl")
sel = json.load(open("selection_frozen.json"))
spec = json.load(open("spec.json"))
ev = {e["arm"]: e for e in C if e["stage"] == "eval"}
seeds = [0, 1, 2]
test_rows = json.load(open("test_ids.json"))  # prompt_id per test row in eval order
tpl = "tinker_cookbook renderer for Qwen/Qwen3.5-2B (renderers.get_recommended_renderer_name), conversation_to_datum(messages, renderer, max_length=2048), TrainOnWhat=ALL_ASSISTANT_MESSAGES default; teacher-forced cross_entropy forward; NLL counted where weight>0"
code = hashlib.sha256(open("run.py", "rb").read() + open("common.py", "rb").read()).hexdigest()
prov = {"code_sha256": code, "prompt_template": tpl, "decoding": {"type": "teacher_forced_forward", "temperature": None, "max_tokens": 0, "stop": []},
        "parser": "token NLL = -sum(logprobs[weights>0]); token_mean = sum NLL / sum counted tokens over 500 items",
        "items": {"ids": [f"no_robots_test:{p}" for p in test_rows]}}
eid = hashlib.sha256(json.dumps(prov, sort_keys=True, separators=(",", ":"), ensure_ascii=True, allow_nan=False).encode()).hexdigest()
arms, loss_eval, scores, obs, preds = {}, {}, [], [], []
for e in (x for x in S if x["stage"] == "eval"):
    arms[e["arm"]] = {"val_token_mean_nll_step40": e["token_mean"]}
    loss_eval[e["arm"]] = {"value": e["token_mean"], "nll_sum": e["nll_sum"], "n_tokens": e["n_tokens"], "n_examples": e["n_examples"], "token_mean": e["token_mean"], "example_mean": e["example_mean"], "aggregation": "token_mean", "n_empty": sum(1 for _, n in e["per_example"] if n == 0)}
diffs = []
for s in seeds:
    b, t = ev.get(f"baseline_adamw_s{s}"), ev.get(f"treatment_dimuon_s{s}")
    for lab, e in (("baseline", b), ("treatment", t)):
        if not e: continue
        arms[e["arm"]] = {"test_token_mean_nll_step100": e["token_mean"], "test_example_mean_nll_step100": e["example_mean"], "lr": next(a for a in json.load(open("confirm_arms.json")) if a["arm"] == e["arm"])["lr"]}
        loss_eval[e["arm"]] = {"value": e["token_mean"], "nll_sum": e["nll_sum"], "n_tokens": e["n_tokens"], "n_examples": e["n_examples"], "token_mean": e["token_mean"], "example_mean": e["example_mean"], "aggregation": "token_mean", "n_empty": sum(1 for _, n in e["per_example"] if n == 0)}
        scores.append({"arm": lab, "seed": s, "metric": "no_robots_test_token_mean_nll_at_step_100", "value": e["token_mean"], "evaluator_id": eid, "predictions_ref": "predictions-confirm-test.jsonl"})
        obs.append({"item_id": "no_robots_test_500_pooled_token_mean", "arm": lab, "seed": s, "metric": "no_robots_test_token_mean_nll_at_step_100", "value": e["token_mean"], "evaluator_id": eid})
        for pid, (nll, n) in zip(test_rows, e["per_example"]):
            preds.append({"item_id": f"no_robots_test:{pid}", "arm": lab, "arm_name": e["arm"], "seed": s, "metric": "example_nll_sum", "prediction": {"nll_sum": nll, "n_tokens": n}, "reference": "no_robots test assistant turns (teacher-forced)", "checkpoint": e["final_state"], "evaluator_id": eid})
    if b and t: diffs.append(b["token_mean"] - t["token_mean"])
res = {"diffs_adamw_minus_dimuon": diffs}
if len(diffs) == 3:
    m, sd = statistics.mean(diffs), statistics.stdev(diffs); hw = 4.302652729911275 * sd / math.sqrt(3)
    res.update(mean=m, sd=sd, lo=m - hw, hi=m + hw)
open("predictions-confirm-test.jsonl", "w").write("".join(json.dumps(p) + "\n" for p in preds))
obsf = "scored-observations-final.jsonl"; open(obsf, "w").write("".join(json.dumps(o) + "\n" for o in obs))
json.dump({"res": res, "arms": arms, "loss_eval": loss_eval, "scores": scores, "eid": eid, "prov": prov,
           "obs_sha": hashlib.sha256(open(obsf, "rb").read()).hexdigest()}, open("final_numbers.json", "w"), indent=1)
print(json.dumps(res), eid)