#!/usr/bin/env python3
"""Grade the Haiku sub-agent battery against the withheld gold key, reusing the
EXACT grader from vibe_eval (math-verify symbolic equiv for non-int golds, strict
int for AIME, letter-match for GPQA). Each agent returned a JSON object mapping
problem id -> final answer in its last assistant message; we extract that JSON
from the agent transcript files without dumping them into anyone's context.
"""
import json, re, sys, glob, os
from collections import defaultdict
from vibe_eval import grade  # reuse identical grading

GOLD = json.load(open("/tmp/haiku_gold.json"))   # id -> [category, answer]
TASKDIR = sys.argv[1]                              # dir holding <agentId>.output files
AGENTS = sys.argv[2:]                              # agentIds to scan

ID_RE = re.compile(r'^(aime|aime25|hmmt25|gpqa)-\d+$')

def extract_preds(path):
    """Find every JSON object in the transcript whose keys look like problem ids;
    merge them (last write wins). Robust to prose around the JSON."""
    raw = open(path, encoding="utf-8", errors="ignore").read()
    preds = {}
    # find balanced-ish {...} blobs; greedy per-line and whole-text candidates
    for m in re.finditer(r'\{[^{}]*\}', raw):
        blob = m.group(0)
        if '"aime' not in blob and '"hmmt' not in blob and '"gpqa' not in blob:
            continue
        try:
            d = json.loads(blob)
        except Exception:
            continue
        for k, v in d.items():
            if ID_RE.match(k) and isinstance(v, str):
                preds[k] = v
    return preds

allpreds = {}
for aid in AGENTS:
    p = os.path.join(TASKDIR, f"{aid}.output")
    if not os.path.exists(p):
        print(f"!! missing {aid}"); continue
    allpreds.update(extract_preds(p))

# grade
cats = defaultdict(lambda: {"n": 0, "ok": 0, "miss": 0})
missing = []
for iid, (cat, gold) in GOLD.items():
    cats[cat]["n"] += 1
    if iid not in allpreds:
        cats[cat]["miss"] += 1; missing.append(iid); continue
    raw = "\\boxed{%s}" % allpreds[iid]      # same wrapping the harness uses for cons grading
    ok, _ = grade(cat, raw, gold)
    cats[cat]["ok"] += int(ok)

order = ["aime", "aime25", "hmmt25", "gpqa"]
print("\n=== HAIKU 4.5 (closed-book, no-tools, sub-agent path) ===")
print(f"{'category':<10}{'pass@1':>9}{'correct':>9}{'n':>5}{'missing':>9}")
for c in order:
    if c not in cats: continue
    v = cats[c]; n = v["n"]
    print(f"{c:<10}{100*v['ok']/n:>8.1f}%{v['ok']:>9}{n:>5}{v['miss']:>9}")
if missing:
    print("\nmissing ids (no answer returned):", missing)
print(f"\ntotal answers parsed: {len(allpreds)}/140")