"""Side-by-side comparison of student_report.py JSONs (e.g. q3 long prompt vs q3s short prompt).

usage: python scripts/compare_reports.py --report long=runs/q3-box-report.json --report short=runs/q3s-box-report.json
           [--extra "short@long_prompt:style_eval=runs/q3s-box-on-long-report.json:style_eval"] --json out.json --md out.md
Rows: every set present in the first report; one column group per report (missing -> "-").
--extra adds a row LABEL taken from SET inside PATH (e.g. the short checkpoint evaluated on long prompts).
"""
import argparse, json
from pathlib import Path

METRICS = [("accuracy", ("accuracy",)), ("record_exact", ("record_exact",)), ("teacher_agreement", ("teacher_agreement",)),
           ("nll", ("nll",)), ("ece", ("ece_15bin",)), ("ece_ts", ("ece_ts_xfit",)),
           ("box_mean_iou", ("box", "overall", "mean_iou")), ("box_iou50", ("box", "overall", "iou50")),
           ("box_iou75", ("box", "overall", "iou75")), ("box_presence", ("box", "overall", "presence_acc")),
           ("box_conf_ece", ("box", "overall", "conf_ece_iou50")), ("sec_per_record", ("mean_sec",))]


def get(d, path):
    for k in path:
        if not isinstance(d, dict) or k not in d:
            return None
        d = d[k]
    return d


def row(report_set):
    a = (report_set or {}).get("all") or {}
    return {m: get(a, p) for m, p in METRICS}


def fmt(v, m):
    if v is None:
        return "-"
    return f"{v:.2f}" if m == "sec_per_record" else f"{v:.4f}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--report", action="append", required=True, help="name=path")
    ap.add_argument("--extra", action="append", default=[], help="label=path:set")
    ap.add_argument("--json"); ap.add_argument("--md")
    a = ap.parse_args()
    reps = [(n, json.loads(Path(p).read_text(encoding="utf-8"))) for n, p in (x.split("=", 1) for x in a.report)]
    out = {"reports": [n for n, _ in reps], "rows": {}}
    for s in reps[0][1]["sets"]:
        out["rows"][s] = {n: row(r["sets"].get(s)) for n, r in reps}
    for x in a.extra:
        label, rest = x.split("=", 1); p, s = rest.rsplit(":", 1)
        out["rows"][label] = {"extra": row(json.loads(Path(p).read_text(encoding="utf-8"))["sets"].get(s))}
    lines = ["| set | metric | " + " | ".join(n for n, _ in reps) + " | delta |", "|---|---|" + "---|" * (len(reps) + 1)]
    for s, cols in out["rows"].items():
        for m, _ in METRICS:
            vals = [cols.get(n, {}).get(m) for n, _ in reps] if "extra" not in cols else [cols["extra"].get(m)]
            d = vals[-1] - vals[0] if len(vals) > 1 and None not in vals else None
            lines.append(f"| {s} | {m} | " + " | ".join(fmt(v, m) for v in vals) + f" | {'-' if d is None else f'{d:+.4f}'} |")
    md = "\n".join(lines)
    print(md)
    if a.json:
        Path(a.json).write_text(json.dumps(out, indent=1), encoding="utf-8")
    if a.md:
        Path(a.md).write_text(md + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()