"""Summarise curve_reeval_<lang>.json (three seeds x rounds 0-8, three runs each) into curve_ungated_<lang>.json for the
3.9 chart: per round, the mean over seeds of tune and unseen scores, the min/max of the per-seed means, the min/max of
single runs, and the mean prompt length (characters for zh; words and characters for en).
Usage: curve_summary.py <lang>"""
import json, re, statistics, sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
lang = sys.argv[1]
re_eval = json.loads((HERE / f"curve_reeval_{lang}.json").read_text(encoding="utf-8"))
words = lambda s: len(re.findall(r"\S+", s))
out = []
for r in range(9):
    rows = [re_eval.get(f"seed{s}:r{r}") for s in (1, 2, 3)]
    rows = [x for x in rows if x]
    if len(rows) < 3:
        break
    row = {"round": r,
           "tune_mean": round(statistics.mean(x["tune_mean"] for x in rows), 3),
           "unseen_mean": round(statistics.mean(x["unseen_mean"] for x in rows), 3),
           "unseen_min": round(min(x["unseen_mean"] for x in rows), 3),
           "unseen_max": round(max(x["unseen_mean"] for x in rows), 3),
           "unseen_run_min": round(min(min(x["unseen"]) for x in rows), 3),
           "unseen_run_max": round(max(max(x["unseen"]) for x in rows), 3),
           "chars_mean": round(statistics.mean(x["chars"] for x in rows))}
    if lang != "zh":
        row["words_mean"] = round(statistics.mean(
            words((HERE / f"runs/{lang}_ungated_{s}/skill_r{r:02d}.md").read_text(encoding="utf-8")) for s in (1, 2, 3)))
    out.append(row)
(HERE / f"curve_ungated_{lang}.json").write_text(json.dumps(out, indent=1), encoding="utf-8")
for row in out:
    print(row)