#!/usr/bin/env python3
"""Aggregate bench_results.jsonl into per-model + per-situation tables."""
import json, os
from collections import defaultdict

BENCH = os.path.expanduser("~/vision_bench")
lines = [json.loads(l) for l in open(f"{BENCH}/bench_results.jsonl")]
recs = [r for r in lines if not r.get("meta")]

models = sorted({r["model"] for r in recs}, key=lambda m: (len(m), m))

print("=" * 100)
print("PER-MODEL SUMMARY")
print("=" * 100)
print(f"{'model':28s} {'wall':>6s} {'eval':>6s} {'tok/s':>7s} {'ident':>5s} {'desc':>5s} {'ground':>6s} {'agent':>5s} {'vram':>6s}")

model_vram = {}
for r in lines:
    if r.get("meta") and r.get("vram"):
        model_vram[r["model"]] = r["vram"].get("vram_gb")

def avg(xs):
    xs = [x for x in xs if x is not None]
    return round(sum(xs)/len(xs), 2) if xs else None

for m in models:
    rs = [r for r in recs if r["model"] == m]
    errs = [r for r in rs if r.get("error")]
    ok = [r for r in rs if not r.get("error")]
    ident = avg([r.get("score") for r in ok if r.get("task") == "identify"])
    desc = avg([r.get("score") for r in ok if r.get("task") == "describe"])
    ground = avg([r.get("score") for r in ok if r.get("task") == "ground"])
    agent = avg([r.get("score") for r in ok if r.get("task") == "agent"])
    wall = avg([r.get("wall_s") for r in ok])
    ev = avg([r.get("eval_s") for r in ok])
    tok = avg([r.get("tok_s") for r in ok])
    vr = model_vram.get(m, "?")
    print(f"{m:28s} {wall:>6} {ev:>6} {tok:>7} {ident:>5} {desc:>5} {ground:>6} {agent:>5} {vr:>6}   (errs: {len(errs)}/{len(rs)})")

print()
print("=" * 100)
print("PER-SITUATION GROUND SCORES (avg across models)")
print("=" * 100)
sits = sorted({r["situation"] for r in recs})
header = f"{'situation':14s}" + "".join(f"{m.split(':')[0][-12:]:>14s}" for m in models)
print(header)
for s in sits:
    row = f"{s:14s}"
    for m in models:
        rs = [r for r in recs if r["model"] == m and r["situation"] == s and not r.get("error")]
        sc = avg([r.get("score") for r in rs if r.get("score") is not None])
        row += f"{sc:>14}"
    print(row)

print()
print("=" * 100)
print("GROUNDING DETAIL (per model, avg distance-ish score per situation)")
print("=" * 100)
for m in models:
    vals = [(r["situation"], r.get("score")) for r in recs if r["model"] == m and r["task"] == "ground" and not r.get("error")]
    if vals:
        print(f"{m:28s}", ", ".join(f"{s[4:]}:{v}" for s, v in vals))