#!/usr/bin/env python3
"""Attempts-to-success distribution for iteration-5 (N=3 variance rerun).

For every CLI verb session across all runs: which attempt number succeeded
(1 / 2 / 3+ / never)? The user-facing goal is "1-2 attempts per call".
"""
import json, glob, os, sys
from pathlib import Path
from collections import defaultdict
sys.path.insert(0, str(Path(__file__).parent))
from analyze_trajectories import load_run, classify_output, cli_verb

def sessions_for(pi):
    calls, meta = load_run(pi)
    cli = []
    for tool, cmd, out in calls:
        if tool != 'bash':
            continue
        cli.append((cli_verb(cmd), cmd, classify_output(out, cmd)))
    # verb -> attempts used until success (per contiguous session)
    stats = defaultdict(lambda: {'sess': 0, 'a1': 0, 'a2': 0, 'a3p': 0, 'never': 0})
    cur = defaultdict(lambda: {'fails': 0, 'done': False})
    for verb, cmd, sig in cli:
        if not verb:
            continue
        st = cur[verb]
        if sig is None:
            if not st['done']:
                stats[verb]['sess'] += 1
                n = st['fails'] + 1
                if n == 1: stats[verb]['a1'] += 1
                elif n == 2: stats[verb]['a2'] += 1
                else: stats[verb]['a3p'] += 1
                st['done'] = True; st['fails'] = 0
        else:
            if not st['done']:
                st['fails'] += 1
    for verb, st in cur.items():
        if not st['done'] and st['fails']:
            stats[verb]['sess'] += 1; stats[verb]['never'] += 1
    return stats

def main():
    roots = sys.argv[1:] or ['/root/eval-rerun/iteration-5']
    cfg_glob = os.environ.get('EVAL_CFG', 'with_skill')
    tot = {'sess': 0, 'a1': 0, 'a2': 0, 'a3p': 0, 'never': 0}
    zero = dict(tot)
    per_verb = defaultdict(lambda: dict(zero))
    npass = nrun = 0
    files = sorted(f for r in roots for f in glob.glob(f'{r}/eval-*/{cfg_glob}/run-1/pi.jsonl'))
    for pi in files:
        nrun += 1
        gp = os.path.join(os.path.dirname(pi), 'grading.json')
        if os.path.exists(gp):
            try:
                g = json.load(open(gp)); exps = g.get('expectations') or []
                if exps and all(x.get('passed') for x in exps):
                    npass += 1
            except Exception:
                pass
        for verb, st in sessions_for(pi).items():
            for k in tot:
                per_verb[verb][k] += st[k]
                tot[k] += st[k]
    print(f"runs={nrun}  eval-pass={npass}/{nrun}")
    s = tot['sess']
    if s:
        print(f"\nverb sessions: {s}")
        print(f"  attempt 1:  {tot['a1']:4}  ({tot['a1']/s*100:.1f}%)")
        print(f"  attempt 2:  {tot['a2']:4}  ({tot['a2']/s*100:.1f}%)")
        print(f"  attempt 3+: {tot['a3p']:4}  ({tot['a3p']/s*100:.1f}%)")
        print(f"  never ok:   {tot['never']:4}  ({tot['never']/s*100:.1f}%)")
        print(f"  ≤2 attempts: {(tot['a1']+tot['a2'])/s*100:.1f}%")
    print(f"\nper-verb (only verbs with a retry or never):")
    for verb, st in sorted(per_verb.items(), key=lambda kv: -(kv[1]['a2']+kv[1]['a3p']+kv[1]['never'])):
        if st['a2'] or st['a3p'] or st['never']:
            print(f"  {verb:24} sess={st['sess']:3}  a1={st['a1']:3} a2={st['a2']} a3+={st['a3p']} never={st['never']}")
    print(f"\nall verbs:")
    for verb, st in sorted(per_verb.items(), key=lambda kv: -kv[1]['sess']):
        print(f"  {verb:24} sess={st['sess']:3}  a1={st['a1']:3} a2={st['a2']} a3+={st['a3p']} never={st['never']}")

if __name__ == '__main__':
    main()