"""Build ``summary.md`` (+ ``summary.json``) from every ``<system>/<set>.preds.jsonl`` under a
bench output dir. Each set gets one table over each system's full sample and, when systems
ran different sample sizes, a second table over the uids every system shares (so the small
frontier sample is compared like for like).

    python -m deg.bench.summary --out /shared/rwkv/results/deg/bench-v1
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from deg.bench import sets as bench_sets
from deg.bench.run import load_preds
from deg.bench.score import aggregate, score_record

COLS = [("valid", "valid_rate", "{:.3f}"), ("decide acc", "decide_acc", "{:.3f}"),
        ("decide ECE", "decide_ece", "{:.3f}"), ("extract EM", "extract_em", "{:.3f}"),
        ("extract norm", "extract_norm", "{:.3f}"), ("count acc", "count_acc", "{:.3f}"),
        ("gen in-bounds", "generate_bounded", "{:.3f}"),
        ("gen ROUGE-L*", "generate_rouge_l", "{:.3f}"),
        ("gen tok/rec", "generated_tokens_mean", "{:.1f}"),
        ("answer tok/rec", "answer_tokens_mean", "{:.1f}"),
        ("prompt tok/rec", "prompt_tokens_mean", "{:.0f}"),
        ("p50 s", "latency_p50_s", "{:.2f}"), ("p95 s", "latency_p95_s", "{:.2f}")]


def _fmt(v: Any, f: str) -> str:
    return "–" if v is None else f.format(v)


def table(rows: list[tuple[str, int, dict[str, Any]]]) -> list[str]:
    head = "| system | n | " + " | ".join(c for c, _, _ in COLS) + " |"
    out = [head, "|" + "---|" * (len(COLS) + 2)]
    for name, n, m in rows:
        out.append(f"| {name} | {n} | " + " | ".join(_fmt(m.get(k), f) for _, k, f in COLS)
                   + " |")
    return out


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=Path("/shared/rwkv/results/deg/bench-v1"))
    ap.add_argument("--records", default=None)
    ap.add_argument("--order", default="deg,qwen3.5,nuextract,clef,frontier",
                    help="system-name prefixes in table order")
    args = ap.parse_args(argv)
    systems = sorted(p for p in args.out.iterdir()
                     if p.is_dir() and p.name not in ("sets", "logs"))
    prefixes = args.order.split(",")
    systems.sort(key=lambda p: next((i for i, x in enumerate(prefixes)
                                     if p.name.startswith(x)), 99))
    set_files = sorted({f.name.removesuffix(".preds.jsonl") for s in systems
                        for f in s.glob("*.preds.jsonl")})
    lines = [f"# deg {args.out.name}: head-to-head", ""]
    summary: dict[str, Any] = {}
    for stem in set_files:
        manifest = args.out / "sets" / f"{stem}.uids"
        name = stem
        recs = {r["uid"]: r for r in bench_sets.sample(name, None, args.out,
                                                       {"records": args.records})}
        order = [u for u in manifest.read_text().split() if u]
        per: dict[str, tuple[Any, dict[str, Any]]] = {}
        for sdir in systems:
            path = sdir / f"{stem}.preds.jsonl"
            if not path.exists():
                continue
            preds = load_preds(path, keep_errors=True)
            mfile = sdir / f"{stem}.metrics.json"
            meta = json.loads(mfile.read_text()) if mfile.exists() else {}
            per[sdir.name] = (preds, meta)
        if not per:
            continue
        summary[stem] = {}
        full_rows, shared = [], None
        for sname, (preds, meta) in per.items():
            uids = [u for u in order if u in preds]
            scope = tuple(meta.get("scope") or ("decide", "extract", "generate"))
            m = aggregate([score_record(recs[u], preds[u], scope) for u in uids],
                          [preds[u] for u in uids])
            full_rows.append((sname, len(uids), m))
            summary[stem][sname] = {"n": len(uids), "scope": list(scope),
                                    "info": meta.get("info"), **m}
            shared = set(uids) if shared is None else shared & set(uids)
        lines += [f"## {name}", "", f"Full sample per system ({len(order)} in manifest).", ""]
        lines += table(full_rows) + [""]
        if shared and any(n != len(shared) for _, n, _ in full_rows):
            uids = [u for u in order if u in shared]
            rows = []
            for sname, (preds, meta) in per.items():
                scope = tuple(meta.get("scope") or ("decide", "extract", "generate"))
                rows.append((sname, len(uids), aggregate(
                    [score_record(recs[u], preds[u], scope) for u in uids],
                    [preds[u] for u in uids])))
            summary[stem]["_shared"] = {s: m for s, _, m in rows}
            lines += [f"Shared subset: the {len(uids)} records every system ran.", ""]
            lines += table(rows) + [""]
    (args.out / "summary.json").write_text(json.dumps(summary, indent=2))
    notes = args.out / "NOTES.md"
    if notes.exists():
        lines += [notes.read_text()]
    (args.out / "summary.md").write_text("\n".join(lines) + "\n")
    print("\n".join(lines))


if __name__ == "__main__":
    main()