"""select: freeze per-rank LR from selection evals. final: paired Hoeffding decision on confirmation evals."""
import json, sys, hashlib, copy, time, os
from statistics import fmean
sys.path.insert(0, __import__("os").path.expanduser("~/repos/research-exchange"))
from research_exchange.research_decisions import assess, assess_spec
from evalgsm import evaluator

LRS = ["1e-4", "3e-4", "1e-3"]
load = lambda f: [json.loads(l) for l in open(f)]
acc = lambda rows: fmean(r["value"] for r in rows)

if sys.argv[1] == "select":
    sel = {"frozen_before_confirmation": True, "frozen_at_unix": time.time(), "search_item_ids": "gsm8k/test/0..199",
           "selection_evaluator_id": evaluator(0, 200)[1], "base_selection_accuracy": acc(load("evals/sel_base.jsonl")), "ranks": {}}
    for R in (1, 32):
        grid = [lr for lr in ["3e-5"] + LRS + ["3e-3"] if os.path.exists(f"evals/sel_rank{R}_lr{lr}.jsonl")]
        scores = {lr: acc(load(f"evals/sel_rank{R}_lr{lr}.jsonl")) for lr in grid}
        best = max(scores.values())
        pick = next(lr for lr in grid if scores[lr] == best)  # grid ascending: ties go to lower LR
        sel["ranks"][str(R)] = {"grid": grid, "extended_with": [lr for lr in grid if lr not in LRS], "selection_accuracy": scores,
                                "selected_lr": pick, "grid_edge": pick in (grid[0], grid[-1])}
    json.dump(sel, open("selection_frozen.json", "w"), indent=1)
    print(json.dumps(sel, indent=1))
else:
    sel = json.load(open("selection_frozen.json"))
    prov, eid = evaluator(200, 1319)
    t = load("evals/conf_treatment.jsonl"); b = load("evals/conf_baseline.jsonl"); base = load("evals/conf_base.jsonl")
    obs = [{k: r[k] for k in ("item_id", "arm", "seed", "metric", "value", "evaluator_id")} for r in t + b]
    with open("scored-observations.jsonl", "w") as fh:
        for o in obs: fh.write(json.dumps(o, sort_keys=True) + "\n")
    digest = hashlib.sha256(open("scored-observations.jsonl", "rb").read()).hexdigest()
    evaluation = {"version": 1, "evaluators": {eid: prov},
                  "scores": [{"arm": a, "metric": "accuracy", "value": acc(rows), "evaluator_id": eid, "seed": 0,
                              "predictions_ref": f"predictions-confirmation-{a}.jsonl"} for a, rows in (("treatment", t), ("baseline", b), ("untrained_base", base))],
                  "comparisons": [{"metric": "accuracy", "arms": ["treatment", "baseline"]}]}
    decision = {"version": 1, "observations": obs, "selection": sel}
    spec = json.load(open("spec.json"))
    rule = copy.deepcopy(spec["decision_rule"])
    rule["adaptive_selection"]["search_item_ids"] = [f"gsm8k/test/{i}" for i in range(200)]
    rule["adaptive_selection"]["confirmation_item_ids"] = [f"gsm8k/test/{i}" for i in range(200, 1319)]
    expanded = assess(rule, decision, evaluation)
    as_saved = assess_spec(spec, decision, evaluation)
    ta, ba, bb = acc(t), acc(b), acc(base)
    out = {"treatment_acc": ta, "baseline_acc": ba, "base_acc": bb, "effect": ta - ba,
           "beats_base_by_0.05": {"treatment": ta - bb >= 0.05, "baseline": ba - bb >= 0.05},
           "grid_edge": any(v["grid_edge"] for v in sel["ranks"].values()),
           "assess_expanded_ids": expanded, "assess_saved_spec": {k: as_saved.get(k) for k in ("status", "issues", "outcome_reason")},
           "observations_sha256": digest, "n_obs": len(obs), "evaluation": evaluation}
    json.dump(out, open("decision_result.json", "w"), indent=1)
    print(json.dumps({k: v for k, v in out.items() if k != "evaluation"}, indent=1, default=str))