#!/usr/bin/env python
"""Collect every finished run into one table, with paired gains against the right baseline.

    python scripts/summarize.py results/runs          # writes results/summary.md and results/summary.json

Run names follow <group>/<model>/<set> for benchmarks (bench/p50c/math500, baseline bench/base/math500)
and dev/<variant> for the practice set (baseline dev/base); quick/<variant> pairs with quick/base the same way.
"""
import json
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp.summary import compare, load_run  # noqa: E402

runs_dir = sys.argv[1] if len(sys.argv) > 1 else os.path.join(ROOT, "results", "runs")
out_dir = os.path.dirname(os.path.abspath(runs_dir))

found = {}
for d, _, files in os.walk(runs_dir):
    if "summary.json" in files and "merged.jsonl" in files:
        found[os.path.relpath(d, runs_dir)] = json.load(open(os.path.join(d, "summary.json")))


def baseline(name):
    """The run to compare against: the unedited model on the same problems, with the same decoding change."""
    p = name.split("/")
    cands = []
    if p[0] == "bench" and len(p) == 3 and p[1] != "base":
        cands = [f"bench/base/{p[2]}"]
    elif p[0] in ("dev", "build", "quick") and len(p) == 2 and p[1] != "base":
        model, _, change = p[1].partition("+")
        cands = ([f"{p[0]}/base+{change}"] if change and model != "base" else []) + [f"{p[0]}/base"]
    return next((c for c in cands if c in found), None)


def pct(x, se=None):
    if x is None:
        return "-"
    return f"{100 * x:.1f}" + (f" ± {100 * se:.1f}" if se is not None else "")


rows, table = [], []
for name in sorted(found):
    s = found[name]
    b = baseline(name)
    c = compare(load_run(os.path.join(runs_dir, b)), load_run(os.path.join(runs_dir, name))) if b else None
    rows.append({"name": name, "summary": {k: v for k, v in s.items() if k not in ("run",)}, "vs": b, "paired": c})
    table.append(f"| {name} | {s['problems']}×{s['run']['n']} | {pct(s['acc'], s['acc_se'])} | "
                 f"{pct(s['cut_share'])} | {s['mean_tokens']:,.0f} | {pct(s['p_correct_given_finished'])} | "
                 f"{pct(s['loop_end_share_of_cut'])} | "
                 + (f"{100 * c['gain']:+.1f} ± {100 * c['gain_se']:.1f} | {c['rescued']} / {c['lost']} |" if c
                    else "- | - |"))
hdr = ("| run | problems × samples | accuracy % | cut at the cap % | mean tokens | correct when finished % "
       "| cut answers ending in repetition % | gain vs baseline (pp) | rescued / lost problems |\n"
       "|---|---|---|---|---|---|---|---|---|")
note = ("\n\"Cut answers ending in repetition\" counts two kinds of ending among the answers cut at the cap: "
        "a word-for-word loop (one block of words repeated back to back at least 3 times, over at least 40 words) "
        "and a looser circling (at least half of the 4-word phrases in the last 200 words occur 3 or more times there). "
        "See verbatim_loop_end and repetitive_end in wsp/grading.py.\n")
md = "# Results\n\n" + hdr + "\n" + "\n".join(table) + "\n" + note
open(os.path.join(out_dir, "summary.md"), "w").write(md)
json.dump(rows, open(os.path.join(out_dir, "summary.json"), "w"), indent=1)
print(md)