"""Download effort-eval results and print accuracy · mean tokens per benchmark, next to base Dory and the
earlier blends on the same problems.

  MODAL_PROFILE=aside python3 scripts/summarize_effort.py smoke
  MODAL_PROFILE=aside python3 scripts/summarize_effort.py full MODEL LOOPS EFFORTS SHARDS
"""
import json
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
RES = ROOT / "data" / "eval" / "results"
BENCHES = ["gpqa", "math500", "aime24", "aime25", "amc23", "heldout"]


def fetch(model, name):
    with tempfile.TemporaryDirectory() as tmp:
        p = subprocess.run(["modal", "volume", "get", "beam-dory-data",
                            f"eval/results/{model.replace('/', '_')}/{name}", tmp], capture_output=True)
        f = Path(tmp) / name
        return [json.loads(l) for l in open(f)] if p.returncode == 0 and f.exists() else None


def row(label, rs, ids=None):
    if ids is not None:
        rs = [r for r in rs if r["id"] in ids]
    cells = []
    for b in BENCHES:
        x = [r for r in rs if r["bench"] == b]
        cells.append(f"{b} {100 * sum(r['correct'] for r in x) / len(x):5.1f}·{sum(r['tokens'] for r in x) / len(x) / 1e3:4.1f}k"
                     if x else f"{b} -")
    tok = sum(r["tokens"] for r in rs) / len(rs)
    acc = 100 * sum(r["correct"] for r in rs) / len(rs)
    print(f"{label:24s} " + "  ".join(cells) + f"  | all {acc:5.1f}·{tok / 1e3:4.2f}k  cut {sum(r['finish'] == 'length' for r in rs)}")
    return rs


def main():
    mode = sys.argv[1]
    if mode == "smoke":
        runs = [(m, 2, e, f"loops2{'' if e == 'default' else '-' + e}-limit.jsonl")
                for m in ("merges/rE-t050", "merges/rE-t030") for e in ("low", "default", "high")]
    else:
        _, _, model, loops, efforts, shards = sys.argv
        runs = [(model, int(loops), e, [f"loops{loops}{'' if e == 'default' else '-' + e}-s{i}of{shards}.jsonl"
                                        for i in range(int(shards))]) for e in efforts.split(",")]
    results = {}
    for model, loops, effort, names in runs:
        names = [names] if isinstance(names, str) else names
        parts = [fetch(model, n) for n in names]
        if all(parts):
            results[(model, loops, effort)] = [r for p in parts for r in p]
            if mode == "full":
                out = RES / f"{model.split('/')[-1]}_loops{loops}-{effort}-64k.jsonl"
                out.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in results[(model, loops, effort)]))
        else:
            print(f"missing: {model} {effort}")
    ids = {r["id"] for rs in results.values() for r in rs}
    loops_set = sorted({k[1] for k in results})
    for loops in loops_set:
        print(f"\n{loops} loops, accuracy · mean output tokens" + (" (same problems)" if mode == "smoke" else ""))
        for ref in ("base", "mix030", "mix050"):
            f = RES / f"{ref}_loops{loops}-64k.jsonl"
            if f.exists():
                row(ref, [json.loads(l) for l in open(f)], ids)
        for (model, l, effort), rs in results.items():
            if l == loops:
                row(f"{model.split('/')[-1]} {effort}", rs)


if __name__ == "__main__":
    main()