import json, subprocess
L = [json.loads(l) for l in open("select_log.jsonl")]
ev = {e["arm"]: e["token_mean"] for e in L if e["stage"] == "eval"}
grid = {"adam": (5e-5, 1e-4, 2e-4), "dimuon": (3e-4, 1e-3, 3e-3)}
sel = {}
for o, lrs in grid.items():
    nm = "adamw" if o == "adam" else o
    scores = {lr: ev[f"sel_{nm}_lr{lr:g}_s0"] for lr in lrs}
    best = min(scores, key=scores.get)
    sel[o] = {"scores": scores, "selected_lr": best, "edge": best in (lrs[0], lrs[-1])}
sel["frozen_before_confirmation"] = True
json.dump(sel, open("selection_frozen.json", "w"), indent=1, default=str); print(json.dumps(sel, default=str))
arms = []
for s in (0, 1, 2):
    arms.append(dict(arm=f"baseline_adamw_s{s}", opt="adam", lr=sel["adam"]["selected_lr"], seed=s, init=f"s{s}"))
    arms.append(dict(arm=f"treatment_dimuon_s{s}", opt="dimuon", lr=sel["dimuon"]["selected_lr"], seed=s, init=f"s{s}"))
json.dump(arms, open("confirm_arms.json", "w"))
subprocess.Popen(f"nohup ../.venv/bin/python -u run.py confirm 100 test '{json.dumps(arms)}' > confirm.out 2>&1 &", shell=True)