"""Merge the 64k continuations (scripts/spawn_continue.py) into full result files, then print accuracy and tokens
at both limits plus paired bootstrap CIs, Dory-Beam-4B minus base, at 64k.

  MODAL_PROFILE=aside python3 scripts/merge_continued.py

Writes data/eval/results/<tag>_[ifeval_]loops<L>-64k.jsonl: the original records, with every answer that hit the
32k limit (16k for IFEval) replaced by its continuation.
"""
import json
import random
import subprocess
import tempfile
from collections import defaultdict
from pathlib import Path

RES = Path(__file__).resolve().parent.parent / "data" / "eval" / "results"
MODELS = {"base": "base", "mix050": "merges_rA-t050"}
BENCHES = ["gpqa", "math500", "aime24", "aime25", "amc23", "heldout"]


def key(r):
    return (r["key"], r["sample"]) if "key" in r else (r["bench"], r["id"], r["sample"])


def merge(tag, loops, kind):
    pre = "ifeval_" if kind == "ifeval" else ""
    with tempfile.TemporaryDirectory() as tmp:
        subprocess.run(["modal", "volume", "get", "beam-dory-data",
                        f"eval/results/{MODELS[tag]}/{pre}loops{loops}-cont64k.jsonl", tmp], check=True,
                       capture_output=True)
        cont = {key(r): r for r in map(json.loads, open(Path(tmp) / f"{pre}loops{loops}-cont64k.jsonl"))}
    orig = [json.loads(l) for l in open(RES / f"{tag}_{pre}loops{loops}.jsonl")]
    assert set(cont) == {key(r) for r in orig if r["finish"] == "length"}, (tag, loops, kind)
    out = [cont.get(key(r), r) for r in orig]
    with open(RES / f"{tag}_{pre}loops{loops}-64k.jsonl", "w") as f:
        for r in out:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    return orig, out


def row(rs, bench):
    rs = [r for r in rs if r["bench"] == bench]
    return (100 * sum(r["correct"] for r in rs) / len(rs), sum(r["tokens"] for r in rs) / len(rs),
            100 * sum(r["finish"] == "length" for r in rs) / len(rs))


def paired_ci(a, b, bench, n_boot=4000):
    pa, pb = defaultdict(list), defaultdict(list)
    for r in a:
        if r["bench"] == bench:
            pa[r["id"]].append(r["correct"])
    for r in b:
        if r["bench"] == bench:
            pb[r["id"]].append(r["correct"])
    d = {i: sum(pa[i]) / len(pa[i]) - sum(pb[i]) / len(pb[i]) for i in pa}
    ids, rng = sorted(d), random.Random(0)
    boots = sorted(sum(d[i] for i in s) / len(s) for s in ([rng.choice(ids) for _ in ids] for _ in range(n_boot)))
    return 100 * sum(d.values()) / len(d), 100 * boots[int(0.025 * n_boot)], 100 * boots[int(0.975 * n_boot)]


def main():
    runs = {}
    for tag in MODELS:
        for loops in (1, 2, 4):
            runs[(tag, loops)] = merge(tag, loops, "bench")
        runs[(tag, "ifeval")] = merge(tag, 2, "ifeval")

    print("accuracy · mean tokens (% cut off), 32k limit -> 64k limit")
    for loops in (1, 2, 4):
        for tag in MODELS:
            o, n = runs[(tag, loops)]
            cells = [f"{b} {row(o, b)[0]:.1f}·{row(o, b)[1] / 1e3:.1f}k({row(o, b)[2]:.0f}%)"
                     f" -> {row(n, b)[0]:.1f}·{row(n, b)[1] / 1e3:.1f}k({row(n, b)[2]:.1f}%)" for b in BENCHES]
            print(f"{tag:7s} L{loops}  " + "  ".join(cells))
    for tag in MODELS:
        o, n = runs[(tag, "ifeval")]
        s = lambda rs: (100 * sum(r["strict_prompt"] for r in rs) / len(rs), sum(r["tokens"] for r in rs) / len(rs),
                        sum(r["finish"] == "length" for r in rs))
        print(f"{tag:7s} IFEval prompt-strict 16k {s(o)[0]:.1f}·{s(o)[1] / 1e3:.1f}k({s(o)[2]} cut)"
              f" -> 64k {s(n)[0]:.1f}·{s(n)[1] / 1e3:.1f}k({s(n)[2]} cut)")

    print("\nDory-Beam-4B minus base at 64k, points, 95% paired bootstrap CI over problems (* = excludes 0)")
    for loops in (1, 2, 4):
        a, b = runs[("mix050", loops)][1], runs[("base", loops)][1]
        cells = []
        for bench in BENCHES:
            d, lo, hi = paired_ci(a, b, bench)
            cells.append(f"{bench} {d:+.1f} [{lo:+.1f},{hi:+.1f}]{'*' if not lo < 0 < hi else ''}")
        ratio = sum(r["tokens"] for r in b) / sum(r["tokens"] for r in a)
        print(f"L{loops}  " + "  ".join(cells) + f"  | base/beam tokens {ratio:.2f}x")
    a, b = runs[("mix050", "ifeval")][1], runs[("base", "ifeval")][1]
    pa = {r["key"]: [] for r in a}
    pb = {r["key"]: [] for r in b}
    for r in a:
        pa[r["key"]].append(r["strict_prompt"])
    for r in b:
        pb[r["key"]].append(r["strict_prompt"])
    d = {k: sum(pa[k]) / len(pa[k]) - sum(pb[k]) / len(pb[k]) for k in pa}
    ids, rng = sorted(d), random.Random(0)
    boots = sorted(sum(d[i] for i in s) / len(s) for s in ([rng.choice(ids) for _ in ids] for _ in range(4000)))
    print(f"IFEval prompt-strict {100 * sum(d.values()) / len(d):+.1f} [{100 * boots[100]:+.1f},{100 * boots[3900]:+.1f}]")


if __name__ == "__main__":
    main()