#!/usr/bin/env python3
"""Aggregate runs.jsonl into summary.json and a check table. No model calls; pure arithmetic over saved rows."""
import collections, json, math, re, sys
from pathlib import Path

BASE = Path(__file__).resolve().parents[1]
ARMS = ['python', 'python+mypy', 'javascript', 'typescript', 'java', 'go', 'rust']


def wilson(k, n, z=1.959963984540054):
    if n == 0:
        return [None, None]
    p = k / n; d = 1 + z * z / n; c = (p + z * z / (2 * n)) / d; h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return [round(c - h, 4), round(c + h, 4)]


def mcnemar_exact(b, c):
    """Two-sided exact binomial test on discordant pairs."""
    n = b + c
    if n == 0:
        return 1.0
    k = min(b, c)
    p = sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n
    return round(min(1.0, 2 * p), 6)


def error_kind(lang, diag):
    pats = {'rs': r'error\[(E\d+)\]', 'ts': r'error (TS\d+)', 'mypy': r'\[([\w-]+)\]\s*$'}
    if lang in ('rs', 'ts'):
        m = re.search(pats[lang], diag)
        return m.group(1) if m else 'syntax/other'
    if lang == 'py':
        m = re.search(pats['mypy'], diag)
        return 'mypy:' + m.group(1) if m else ('SyntaxError' if 'Syntax' in diag or 'Indentation' in diag else 'other')
    if lang == 'go':
        for key in ['undefined', 'declared and not used', 'imported and not used', 'cannot use', 'mismatched types', 'missing return', 'syntax error', 'expected']:
            if key in diag:
                return key
        return 'other'
    if lang == 'java':
        for key in ['cannot find symbol', 'incompatible types', 'missing return', 'possible lossy conversion', 'bad operand', 'expected', 'unreachable', 'might not have been initialized']:
            if key in diag:
                return key
        return 'other'
    return 'SyntaxError' if 'Syntax' in diag else 'other'


STD = r'undefined: (strings|sort|math|fmt|strconv|unicode|bytes|regexp|reflect)\b'


def rejection_class(lang, diag):
    """syntax: malformed text; import: a standard-library package used without importing it;
    semantic: everything else (types, names, ownership, missing returns) — what a type system is for."""
    d = diag
    if lang == 'go' and re.search(STD, d):
        return 'import'
    if lang == 'py' and re.search(r'Name "(math|re|collections|itertools|functools|string|heapq|bisect)" is not defined', d):
        return 'import'
    syntax = {
        'py': lambda: not re.search(r':\d+: error:', d),  # anything before mypy is compile()'s syntax check
        'js': lambda: True,  # node --check is syntax-only
        'ts': lambda: bool(re.search(r'error TS1\d{3}', d)),
        'java': lambda: bool(re.search(r"illegal start|expected$|';' expected|unclosed|reached end of file|not a statement", d)),
        'go': lambda: bool(re.search(r'syntax error|expected', d)),
        'rs': lambda: bool(re.search(r'error: (expected|unexpected|mismatched closing|this file contains an unclosed)', d)),
    }[lang]()
    if syntax:
        return 'syntax'
    return 'semantic'


def arm_stats(a, rs):
    n = len(rs)
    f_pass = sum(r['first_pass'] for r in rs)
    f_gate = sum(r['first_gate_ok'] for r in rs)
    fin = sum(r['final_pass'] for r in rs)
    fin_gate = sum(r['final_gate_ok'] for r in rs)
    wrong0 = [r for r in rs if not r['first_pass']]
    caught = sum(not r['first_gate_ok'] for r in wrong0)
    rejected0 = [r for r in rs if not r['first_gate_ok']]
    false_alarm = sum(r['first_pass'] for r in rejected0)
    silent = sum(r['final_gate_ok'] and not r['final_pass'] for r in rs)
    gained = sum((not r['first_pass']) and r['final_pass'] for r in rs)
    lost = sum(r['first_pass'] and not r['final_pass'] for r in rs)
    kinds = collections.Counter(); classes = collections.Counter(); sem_caught = 0
    for r in rejected0:
        d = r['attempts'][0]['diagnostics'][0]
        kinds[error_kind(r['lang'], d)] += 1
        cls = rejection_class(r['lang'], d)
        classes[cls] += 1
        sem_caught += cls == 'semantic' and not r['first_pass']
    tokens = sum(r['tokens'] for r in rs)
    # "Shipped and correct": the gate accepted the program AND the hidden tests pass. A program the
    # gate rejects is never shipped in a compiler-in-the-loop workflow, even if its emitted JS would pass.
    ship0 = sum(r['first_gate_ok'] and r['first_pass'] for r in rs)
    ship = sum(r['final_gate_ok'] and r['final_pass'] for r in rs)
    ship_first = sum(r['final_gate_ok'] and r['final_pass'] and r['repairs'] == 0 for r in rs)
    gain_by_class = collections.Counter(rejection_class(r['lang'], r['attempts'][0]['diagnostics'][0]) for r in rs
                                        if not r['first_gate_ok'] and r['final_gate_ok'] and r['final_pass'])
    return {
        'first_pass': f_pass, 'first_pass_rate': round(f_pass / n, 4), 'first_pass_ci95': wilson(f_pass, n),
        'first_gate_reject': n - f_gate, 'final_pass': fin, 'final_pass_rate': round(fin / n, 4), 'final_pass_ci95': wilson(fin, n),
        'gained_by_feedback': gained, 'lost_by_feedback': lost, 'mcnemar_p_first_vs_final': mcnemar_exact(gained, lost),
        'first_wrong': len(wrong0), 'first_wrong_caught_by_gate': caught,
        'catch_rate': round(caught / len(wrong0), 4) if wrong0 else None, 'catch_rate_ci95': wilson(caught, len(wrong0)),
        'gate_rejects_that_would_pass': false_alarm, 'gate_runnable_when_rejected': a.split('-')[0] in ('python', 'python+mypy', 'javascript', 'typescript'),
        'final_gate_ok': fin_gate, 'final_silent_wrong': silent, 'final_still_rejected': n - fin_gate,
        'pass_given_gate_ok': round(fin / fin_gate, 4) if fin_gate else None,
        'repair_rounds': sum(r['repairs'] for r in rs), 'extract_failures': sum(a2.get('extract_failed', False) for r in rs for a2 in r['attempts']),
        'output_tokens': tokens, 'tokens_per_final_pass': round(tokens / fin, 1) if fin else None,
        'first_rejection_kinds': dict(kinds.most_common()),
        'first_rejection_classes': dict(classes), 'first_wrong_caught_semantic': sem_caught,
        'semantic_catch_rate': round(sem_caught / len(wrong0), 4) if wrong0 else None, 'semantic_catch_ci95': wilson(sem_caught, len(wrong0)),
        'first_attempt_tokens': sum(r['attempts'][0]['eval_count'] for r in rs),
        'shipped_correct_no_feedback': ship0, 'shipped_correct_no_feedback_ci95': wilson(ship0, n),
        'shipped_correct': ship, 'shipped_correct_rate': round(ship / n, 4), 'shipped_correct_ci95': wilson(ship, n),
        'shipped_correct_first_try': ship_first, 'feedback_gains_by_first_rejection_class': dict(gain_by_class),
        'tokens_per_shipped_correct': round(tokens / ship, 1) if ship else None, 'shipped_correct_after_repair': ship - ship_first,
        'rejected_but_tests_pass_final': sum((not r['final_gate_ok']) and r['final_pass'] for r in rs),
    }


def main(path=BASE / 'results/runs.jsonl'):
    rows = [json.loads(x) for x in Path(path).read_text().splitlines()]
    by = collections.defaultdict(dict)
    for r in rows:
        by[r['arm']][r['task']] = r
    common = sorted(set.intersection(*[set(by[a]) for a in ARMS]))
    n = len(common)
    out = {'tasks': n, 'task_ids': common, 'model': rows[0]['model'], 'arms': {}}
    for a in ARMS:
        out['arms'][a] = arm_stats(a, [by[a][t] for t in common])
    # Paired comparisons of final pass between arms on identical tasks.
    pairs = {}
    for x, y in [('python', 'python+mypy'), ('javascript', 'typescript'), ('python', 'rust'), ('python', 'typescript'), ('python', 'go'), ('python', 'java'), ('typescript', 'rust')]:
        ok = lambda r: r['final_gate_ok'] and r['final_pass']
        b = sum(ok(by[x][t]) and not ok(by[y][t]) for t in common)
        c = sum(ok(by[y][t]) and not ok(by[x][t]) for t in common)
        pairs[f'{x} vs {y}'] = {'only_first_passes': b, 'only_second_passes': c, 'mcnemar_p': mcnemar_exact(b, c)}
    out['paired_final'] = pairs
    allp = []
    for i, x in enumerate(ARMS):
        for y in ARMS[i + 1:]:
            ok = lambda r: r['final_gate_ok'] and r['final_pass']
            b = sum(ok(by[x][t]) and not ok(by[y][t]) for t in common); c = sum(ok(by[y][t]) and not ok(by[x][t]) for t in common)
            allp.append([x, y, b, c, mcnemar_exact(b, c)])
    m = len(allp); order = sorted(range(m), key=lambda k: allp[k][4]); running = 0
    for rank, k in enumerate(order):
        running = max(running, min(1.0, (m - rank) * allp[k][4])); allp[k].append(round(running, 6))
    out['all_pairs_holm'] = [dict(zip(['a', 'b', 'only_a', 'only_b', 'p', 'p_holm'], v)) for v in allp]
    # Sensitivity arm: TypeScript re-run with JavaScript's stop list (adds "\nconsole.log"), same tasks and settings.
    sp = BASE / 'results/sensitivity-ts-jsstops.jsonl'
    if sp.exists():
        srows = {r['task']: r for r in map(json.loads, sp.read_text().splitlines())}
        if set(common) <= set(srows):
            sa = 'typescript-jsstops'; by[sa] = srows
            sens = {'arm': sa, 'stats': arm_stats(sa, [srows[t] for t in common]),
                    'first_completions_with_console_log': sum('console.log' in srows[t]['attempts'][0]['completion'] for t in common),
                    'original_ts_first_completions_with_console_log': sum('console.log' in by['typescript'][t]['attempts'][0]['completion'] for t in common), 'paired': {}}
            for other in ('javascript', 'typescript'):
                for stage, f in (('no_feedback', lambda r: r['first_gate_ok'] and r['first_pass']), ('with_feedback', lambda r: r['final_gate_ok'] and r['final_pass'])):
                    b = sum(f(by[other][t]) and not f(srows[t]) for t in common); c = sum(f(srows[t]) and not f(by[other][t]) for t in common)
                    sens['paired'][f'{other} vs {sa} ({stage})'] = {'only_first_passes': b, 'only_second_passes': c, 'mcnemar_p': mcnemar_exact(b, c)}
            out['sensitivity_ts_jsstops'] = sens
    (BASE / 'results/summary.json').write_text(json.dumps(out, indent=2) + '\n')
    lines0 = ['| Arm | Shipped & correct, no feedback | Shipped & correct, with ≤2 compiler-feedback rounds | 95% CI |', '|---|---:|---:|---|']
    for a in ARMS:
        s_ = out['arms'][a]
        lines0.append(f"| {a} | {s_['shipped_correct_no_feedback']}/{n} | {s_['shipped_correct']}/{n} ({s_['shipped_correct_rate']:.1%}) | {s_['shipped_correct_ci95'][0]:.1%}–{s_['shipped_correct_ci95'][1]:.1%} |")
    hdr = '| Arm | First-try pass | Rejected by gate (first try) | Wrong first tries caught by gate | Final pass after ≤2 compiler-feedback rounds | Gained / lost | Compiles but wrong (final) | Pass rate when gate is green |'
    lines = [hdr, '|---|---:|---:|---:|---:|---:|---:|---:|']
    for a in ARMS:
        s = out['arms'][a]
        lines.append(f"| {a} | {s['first_pass']}/{n} ({s['first_pass_rate']:.1%}) | {s['first_gate_reject']} | {s['first_wrong_caught_by_gate']}/{s['first_wrong']} ({(s['catch_rate'] or 0):.0%}) | "
                     f"{s['final_pass']}/{n} ({s['final_pass_rate']:.1%}) | +{s['gained_by_feedback']} / −{s['lost_by_feedback']} | {s['final_silent_wrong']} | {(s['pass_given_gate_ok'] or 0):.1%} |")
    lines = lines0 + [''] + lines
    (BASE / 'results/table.md').write_text('\n'.join(lines) + '\n')
    print('\n'.join(lines))
    print(json.dumps(pairs, indent=1))
    for a in ARMS:
        print(a, out['arms'][a]['first_rejection_kinds'], 'false alarms:', out['arms'][a]['gate_rejects_that_would_pass'], 'tokens/pass', out['arms'][a]['tokens_per_final_pass'])


if __name__ == '__main__':
    main(*sys.argv[1:])