#!/usr/bin/env python3
"""Count graded episodes in a run dir. `eval --resume` appends to the same traces.jsonl, so
log-line counts undercount across resumes — this is the number that actually matters."""
import json, os, sys

# `-n` prints just the graded count, one per line. Needed because a naive
# `grep -o 'graded= *[0-9]*'` also matches the "ungraded=" field and returns two numbers.
QUIET = "-n" in sys.argv
for d in [a for a in sys.argv[1:] if a != "-n"]:
    p = os.path.join(d, "traces.jsonl")
    graded = ungraded = 0
    if os.path.exists(p):
        for line in open(p):
            for t in json.loads(line).get("traces", []):
                if t.get("rewards"):
                    graded += 1
                else:
                    ungraded += 1
    if QUIET:
        print(graded)
    else:
        print(f"{os.path.basename(d):18s} graded={graded:4d} ungraded={ungraded:4d}")