"""Summarize runs: final/min val loss per run, plus tokens needed to reach target losses."""

import argparse
import json
from pathlib import Path


def load(run_dir):
    recs = [json.loads(l) for l in (run_dir / "metrics.jsonl").read_text().splitlines() if l.strip()]
    val = [(r["tokens"], r["val_loss"]) for r in recs if "val_loss" in r]
    return val


def main():
    p = argparse.ArgumentParser()
    p.add_argument("runs_dir")
    p.add_argument("--targets", default="4.0,3.6,3.4")
    args = p.parse_args()
    targets = [float(t) for t in args.targets.split(",")]
    rows = []
    for d in sorted(Path(args.runs_dir).iterdir()):
        if not (d / "metrics.jsonl").exists():
            continue
        val = load(d)
        if not val:
            continue
        reach = [next((t for t, v in val if v <= tg), None) for tg in targets]
        rows.append((d.name, "done" if (d / "DONE").exists() else "running", val[0][1], val[-1][0], val[-1][1], reach))
    hdr = f"{'run':40s} {'status':8s} {'init_loss':>9s} {'tokens':>8s} {'val':>7s} " + " ".join(f"{'@' + str(t):>8s}" for t in targets)
    print(hdr)
    for name, st, v0, tok, v, reach in rows:
        r = " ".join(f"{(f'{x / 1e6:.0f}M' if x else '-'):>8s}" for x in reach)
        print(f"{name:40s} {st:8s} {v0:9.3f} {tok / 1e6:7.0f}M {v:7.3f} {r}")


if __name__ == "__main__":
    main()