"""2.8 (English edition): grade every raw output and aggregate.

  grade.py              judge everything in raw/ not yet judged with the primary judge (lib.PRIMARY_JUDGE: gemini grades
                        deepseek and grok, deepseek grades gemini; cached in graded/), then write results.json
  grade.py secondary    second opinion on every answer from the other judge model (deepseek for deepseek/grok answers,
                        gemini for gemini answers), to measure judge-judge agreement; stored as correct_secondary
  grade.py sample N     write spotcheck.json: N primary judgments sampled with a fixed seed for hand review
  grade.py agree        compare spotcheck_review.json (hand labels) with the judge
  grade.py analyze      extra tables for the README (index_fixed vs control, cue split, failure taxonomy)

An answer is flagged web=true when the gateway orchestrator ran a web search/extract during that task."""
import hashlib, json, random, sys
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from lib import HERE, judge, load_pool, PRIMARY_JUDGE, append_costs

RAW, GRADED = HERE / 'raw', HERE / 'graded'
GRADED.mkdir(exist_ok=True)
TASKS = {t['id']: t for t in json.loads((HERE / 'tasks.json').read_text())}
ORDER = ['none', 'full_1', 'full_10', 'full_30', 'full_all', 'index', 'index_rep2', 'index_fixed']
SECONDARY_JUDGE = {'deepseek-v4.1-flash': 'deepseek-v4.1-flash', 'grok-4-fast': 'deepseek-v4.1-flash',
                   'gemini-3.8-flash': 'gemini-3.8-flash'}


def h(s):
    return hashlib.sha1((s or '').encode()).hexdigest()[:12]


def grade_all(secondary=False):
    key = 'correct_secondary' if secondary else 'correct'
    jobs = []
    for f in sorted(RAW.glob('*.json')):
        gpath = GRADED / f.name
        cache = {d['id']: d for d in json.loads(gpath.read_text())} if gpath.exists() else {}
        for d in json.loads(f.read_text()):
            c = cache.get(d['id'])
            if c and c.get('answer_hash') == h(d.get('answer')) and c.get(key) is not None:
                continue
            jobs.append((f.name, d))

    def one(job):
        name, d = job
        jm = (SECONDARY_JUDGE if secondary else PRIMARY_JUDGE)[d['model']]
        if d.get('error') or d.get('answer') is None:
            g = {'correct': False, 'reason': f'no answer: {d.get("error")}', 'judge': jm, 'judge_cost': 0}
        else:
            g = judge(TASKS[d['id']], d['answer'], jm)
        return name, d, g

    print('judging', len(jobs), 'secondary' if secondary else 'primary', flush=True)
    store, usd = {}, 0.0
    with ThreadPoolExecutor(12) as ex:
        for i, (name, d, g) in enumerate(ex.map(one, jobs)):
            store.setdefault(name, []).append((d, g))
            usd += g.get('judge_cost') or 0
            if i % 50 == 0:
                print(i, f'${usd:.3f}', flush=True)
    for name, items in store.items():
        gpath = GRADED / name
        cache = {d['id']: d for d in json.loads(gpath.read_text())} if gpath.exists() else {}
        for d, g in items:
            c = cache.get(d['id'], {'id': d['id']})
            if c.get('answer_hash') != h(d.get('answer')):
                c = {'id': d['id']}
            c['answer_hash'] = h(d.get('answer'))
            if secondary:
                c.update({'correct_secondary': g['correct'], 'reason_secondary': g['reason'], 'judge_secondary': g.get('judge')})
            else:
                c.update({'correct': g['correct'], 'reason': g['reason'], 'judge': g.get('judge')})
            cache[d['id']] = c
        gpath.write_text(json.dumps(sorted(cache.values(), key=lambda x: x['id']), ensure_ascii=False, indent=1))
    if jobs:
        append_costs(f'2_8_skills/grade.py{" secondary" if secondary else ""}', len(jobs), 0, 0, usd,
                     'judge calls; token counts not tracked here, $ from gateway trace')
    aggregate()


def load_joined():
    rows = []
    for f in sorted(RAW.glob('*.json')):
        gpath = GRADED / f.name
        gs = {d['id']: d for d in json.loads(gpath.read_text())} if gpath.exists() else {}
        for d in json.loads(f.read_text()):
            g = gs.get(d['id'], {})
            ok_hash = g.get('answer_hash') == h(d.get('answer'))
            rows.append({**d, 'correct': g.get('correct') if ok_hash else None, 'judge_reason': g.get('reason'),
                         'judge': g.get('judge'), 'correct_secondary': g.get('correct_secondary') if ok_hash else None,
                         'web': bool(d.get('web_tools'))})
    return rows


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


def aggregate():
    pool, _ = load_pool()
    rows = load_joined()
    res = {'pool': {'n_skills': pool['n_skills'], 'total_tokens_cl100k': pool['total_tokens_cl100k']}, 'cells': {}}
    for model in sorted({r['model'] for r in rows}):
        for cond in ORDER:
            rs = [r for r in rows if r['model'] == model and r['condition'] == cond]
            if not rs:
                continue
            ok = [r for r in rs if not r.get('error') and r.get('answer') is not None]
            cell = {'n': len(rs), 'n_errors': len(rs) - len(ok),
                    'errors_sample': sorted({(r.get('error') or '')[:160] for r in rs if r.get('error')})[:3],
                    'correct': sum(1 for r in rs if r['correct']),
                    'ungraded': sum(1 for r in rs if r['correct'] is None),
                    'judge': sorted({r.get('judge') for r in rs if r.get('judge')}),
                    'correct_secondary_judge': sum(1 for r in rs if r.get('correct_secondary')),
                    'n_secondary_graded': sum(1 for r in rs if r.get('correct_secondary') is not None),
                    'judges_disagree': sum(1 for r in rs if r.get('correct_secondary') is not None
                                           and r['correct'] is not None and r['correct'] != r['correct_secondary']),
                    'mean_prompt_tokens': mean([r.get('prompt_tokens') for r in ok]),
                    'mean_completion_tokens': mean([r.get('completion_tokens') for r in ok]),
                    'mean_latency_s': mean([r.get('sec') for r in ok]),
                    'total_cost_usd': sum(r.get('cost') or 0 for r in rs),
                    'mean_cost_usd': mean([r.get('cost') for r in ok]),
                    'n_web_flagged': sum(r['web'] for r in rs)}
            cell['accuracy'] = cell['correct'] / len(rs)
            clean = [r for r in rs if not r['web']]
            cell['correct_excl_web'] = sum(1 for r in clean if r['correct'])
            cell['n_excl_web'] = len(clean)
            if cond.startswith('index'):
                cell['opened_target'] = sum(1 for r in rs if r.get('opened_target'))
                cell['mean_opens'] = mean([len(r.get('opened') or []) for r in rs])
                cell['opened_nothing'] = sum(1 for r in rs if not r.get('requested'))
                wrong = [r for r in rs if not r['correct']]
                cell['wrong_breakdown'] = {
                    'no_skill_opened': sum(1 for r in wrong if not r.get('requested')),
                    'only_other_skills_opened': sum(1 for r in wrong if r.get('requested') and not r.get('opened_target')),
                    'target_opened_but_wrong': sum(1 for r in wrong if r.get('opened_target')),
                    'error_no_answer': sum(1 for r in wrong if r.get('error'))}
                cell['correct_given_target_opened'] = f"{sum(1 for r in rs if r.get('opened_target') and r['correct'])}/{cell['opened_target']}"
                cell['right_without_target'] = sum(1 for r in rs if r['correct'] and not r.get('opened_target'))
                cell['other_opened'] = Counter(n for r in rs for n in (r.get('opened') or []) if n != r['target']).most_common(15)
            res['cells'][f'{model}:{cond}'] = cell
    old = json.loads((HERE / 'results.json').read_text()) if (HERE / 'results.json').exists() else {}
    for k in ('analysis', 'spotcheck', 'meta'):
        if k in old:
            res[k] = old[k]
    (HERE / 'results.json').write_text(json.dumps(res, ensure_ascii=False, indent=1))
    print(f'{"cell":36s} {"acc":>7s} {"err":>3s} {"web":>3s} {"tok":>9s} {"lat":>6s} {"cost$":>8s} {"j2":>5s}')
    for k, c in res['cells'].items():
        print(f'{k:36s} {c["correct"]:>3d}/{c["n"]:<3d} {c["n_errors"]:>3d} {c["n_web_flagged"]:>3d} '
              f'{(c["mean_prompt_tokens"] or 0):>9.0f} {(c["mean_latency_s"] or 0):>6.1f} {c["total_cost_usd"]:>8.3f} '
              f'{c["correct_secondary_judge"]:>2d}/{c["n_secondary_graded"]:<2d}'
              + (f'  opened_target={c["opened_target"]} wrong={c["wrong_breakdown"]}' if 'opened_target' in c else ''))


def sample(n):
    rows = [r for r in load_joined() if r.get('answer')]
    rng = random.Random(7)
    pos = [r for r in rows if r['correct']]
    neg = [r for r in rows if r['correct'] is False]
    pick = rng.sample(pos, min(n // 2, len(pos))) + rng.sample(neg, min(n - n // 2, len(neg)))
    rng.shuffle(pick)
    out = [{'model': r['model'], 'condition': r['condition'], 'id': r['id'], 'question': TASKS[r['id']]['question'],
            'excerpt': TASKS[r['id']]['excerpt'], 'rubric': TASKS[r['id']]['rubric'], 'answer': r['answer'],
            'judge': r.get('judge'), 'judge_correct': r['correct'], 'judge_reason': r['judge_reason']} for r in pick]
    (HERE / 'spotcheck.json').write_text(json.dumps(out, ensure_ascii=False, indent=1))
    print('wrote', len(out))


def agree():
    sc = json.loads((HERE / 'spotcheck.json').read_text())
    rv = json.loads((HERE / 'spotcheck_review.json').read_text())
    n = len(rv)
    same = sum(1 for s, v in zip(sc, rv) if s['judge_correct'] == v['my_correct'])
    print(f'agreement {same}/{n}')
    dis = []
    for s, v in zip(sc, rv):
        if s['judge_correct'] != v['my_correct']:
            dis.append({'model': s['model'], 'condition': s['condition'], 'id': s['id'], 'judge': s['judge_correct'],
                        'me': v['my_correct'], 'note': v.get('note')})
            print('DISAGREE', dis[-1])
    res = json.loads((HERE / 'results.json').read_text())
    res['spotcheck'] = {'n': n, 'agree': same, 'disagreements': dis}
    (HERE / 'results.json').write_text(json.dumps(res, ensure_ascii=False, indent=1))


def analyze():
    rows = load_joined()
    by = {(r['model'], r['condition'], r['id']): r for r in rows}
    models = sorted({r['model'] for r in rows})
    no_cue = set(json.loads((HERE / 'task_cues.json').read_text())['no_cue'])
    plan_path = HERE / 'index_fixed_plan.json'
    plan = json.loads(plan_path.read_text()) if plan_path.exists() else {}
    out = {}
    for m in models:
        a = {}
        p = plan.get(m)
        if p and (m, 'index_fixed', (p['failed_not_opened'] + p['passing_sample'] or ['-'])[0]) in by:
            f = p['failed_not_opened']
            g = lambda c, i: by.get((m, c, i), {})
            a['index_fixed'] = {
                'failed_not_opened': len(f),
                'fixed_now_correct': [i for i in f if g('index_fixed', i).get('correct')],
                'fixed_now_opened_target': [i for i in f if g('index_fixed', i).get('opened_target')],
                'fixed_still_opened_nothing': [i for i in f if not g('index_fixed', i).get('requested')],
                'control_rep2_now_correct': [i for i in f if g('index_rep2', i).get('correct')],
                'control_rep2_now_opened_target': [i for i in f if g('index_rep2', i).get('opened_target')],
                'passing_sample': len(p['passing_sample']),
                'passing_sample_still_correct_fixed': sum(1 for i in p['passing_sample'] if g('index_fixed', i).get('correct')),
                'passing_sample_still_correct_rep2': sum(1 for i in p['passing_sample'] if g('index_rep2', i).get('correct')),
                'passing_sample_broke_fixed': [i for i in p['passing_sample'] if not g('index_fixed', i).get('correct')]}
        for cond in ('index', 'index_rep2'):
            idx = [by[(m, cond, t)] for t in TASKS if (m, cond, t) in by]
            for lab, sel in (('cue', lambda r: r['id'] not in no_cue), ('no_cue', lambda r: r['id'] in no_cue)):
                rs = [r for r in idx if sel(r)]
                a[f'{cond}_{lab}'] = {'n': len(rs), 'correct': sum(1 for r in rs if r['correct']),
                                      'opened_target': sum(1 for r in rs if r.get('opened_target')),
                                      'opened_nothing': sum(1 for r in rs if not r.get('requested'))}
        idx = [by[(m, 'index', t)] for t in TASKS if (m, 'index', t) in by]
        rep = [by[(m, 'index_rep2', t)] for t in TASKS if (m, 'index_rep2', t) in by]
        if idx and rep:
            a['index_vs_rep2_flips'] = [r['id'] for r in idx if by[(m, 'index_rep2', r['id'])]['correct'] != r['correct']]
        ok1 = [t for t in TASKS if (m, 'full_1', t) in by and by[(m, 'full_1', t)]['correct']]
        a['full_1_correct'] = len(ok1)
        a['dilution'] = {}
        for k in ('full_10', 'full_30', 'full_all'):
            if all((m, k, t) in by for t in TASKS) and not by[(m, k, 'T01')].get('error'):
                lost = [t for t in ok1 if not by[(m, k, t)]['correct']]
                gained = [t for t in TASKS if (m, 'full_1', t) in by and not by[(m, 'full_1', t)]['correct'] and by[(m, k, t)]['correct']]
                a['dilution'][k] = {'lost_vs_full_1': lost, 'gained_vs_full_1': gained,
                                    'lost_target_pos': {t: by[(m, k, t)].get('target_pos') for t in lost}}
        a['index_right_without_target'] = [r['id'] for r in idx if r['correct'] and not r.get('opened_target')]
        a['index_opened_counts'] = dict(Counter(len(r.get('requested') or []) for r in idx))
        out[m] = a
    res = json.loads((HERE / 'results.json').read_text())
    res['analysis'] = out
    (HERE / 'results.json').write_text(json.dumps(res, ensure_ascii=False, indent=1))
    print(json.dumps(out, ensure_ascii=False, indent=1))


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    if cmd == 'analyze':
        analyze()
    elif cmd == 'sample':
        sample(int(sys.argv[2]))
    elif cmd == 'agree':
        agree()
    elif cmd == 'agg':
        aggregate()
    elif cmd == 'secondary':
        grade_all(secondary=True)
    else:
        grade_all()