"""Paired comparison between eval rows.

Every row in ``results.jsonl`` uses the same held-out start/goal pairs, so
success rates must be compared per-episode rather than as two independent
binomials. With 50 episodes one flipped episode moves the rate by 2 points,
which is well inside the noise a naive comparison would ignore.

McNemar's exact test uses only the discordant episodes (the ones where the
two planners disagree), which is the right conditioning for paired binary
outcomes. The bootstrap CI resamples episodes to put an interval on the
difference.
"""

import argparse
import json
from math import comb
from pathlib import Path

import numpy as np


def mcnemar_exact(a, b):
    """Two-sided exact McNemar p-value for paired binary outcomes."""
    only_a = int(np.sum(a & ~b))
    only_b = int(np.sum(b & ~a))
    n = only_a + only_b
    if n == 0:
        return 1.0, only_a, only_b
    k = min(only_a, only_b)
    tail = sum(comb(n, i) for i in range(k + 1)) / 2**n
    return min(1.0, 2 * tail), only_a, only_b


def bootstrap_ci(a, b, reps=20000, seed=0):
    """Percentile CI on the paired success-rate difference, in points."""
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, len(a), size=(reps, len(a)))
    diffs = (a[idx].mean(axis=1) - b[idx].mean(axis=1)) * 100
    return np.percentile(diffs, [2.5, 97.5])


def label(row):
    tag = row['planner']
    if row.get('receding_horizon', 1) != 1:
        tag += f'+exec{row["receding_horizon"]}'
    ckpt = row.get('checkpoint') or ''
    # exp7 cell names first (fused192_s2/base_r2 before their shorter
    # prefixes so a checkpoint doesn't also pick up an unrelated substring)
    for name in ('fused192_s2', 'fused192', 'fused256', 'w192np_split',
                 'base_r2', 'terminal_only', 'no_support', 'ah_hold0.5',
                 'ah_hold0.0', 'ah_hold1.0'):
        if name in ckpt:
            tag += f'[{name}]'
            break
    return tag


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--results', default='data/runs/eval/results.jsonl')
    p.add_argument('--pairs', nargs='*', default=None,
                   help='LABEL_A vs LABEL_B, as "A::B" strings')
    args = p.parse_args()

    rows = [
        json.loads(x)
        for x in Path(args.results).read_text().splitlines()
        if x.strip()
    ]
    rows = [r for r in rows if r.get('episode_successes')]
    if not rows:
        print('no rows carry per-episode outcomes yet — rerun evals')
        return

    # Pool across eval seeds: multiple rows can share one label (same
    # checkpoint + rh, different --seed), and the pairing is valid *within*
    # each seed (identical held-out start/goal pairs across planners at that
    # seed) — so seed order must match between any two labels being compared.
    # A plain dict keyed by label would instead let the last seed's row
    # silently overwrite the earlier ones.
    grouped = {}
    for r in rows:
        grouped.setdefault(label(r), []).append(
            (r.get('seed', 0), np.array(r['episode_successes'], dtype=bool))
        )
    by_label = {}
    for lbl, items in grouped.items():
        items.sort(key=lambda item: item[0])
        by_label[lbl] = np.concatenate([arr for _, arr in items])

    print('available rows with paired outcomes:')
    for k, v in by_label.items():
        print(f'  {k:<40} {v.mean() * 100:5.1f}%  (n={len(v)})')

    pairs = []
    if args.pairs:
        for spec in args.pairs:
            a, b = spec.split('::')
            pairs.append((a, b))
    else:
        keys = list(by_label)
        pairs = [(keys[i], keys[j])
                 for i in range(len(keys)) for j in range(i + 1, len(keys))]

    print(f'\n{"comparison":<62}{"diff":>8}{"95% CI":>18}{"p":>9}')
    print('-' * 97)
    for a, b in pairs:
        if a not in by_label or b not in by_label:
            print(f'  skip {a} vs {b}: missing')
            continue
        sa, sb = by_label[a], by_label[b]
        if len(sa) != len(sb):
            print(f'  skip {a} vs {b}: different episode counts')
            continue
        diff = (sa.mean() - sb.mean()) * 100
        lo, hi = bootstrap_ci(sa, sb)
        pval, na, nb = mcnemar_exact(sa, sb)
        star = '*' if pval < 0.05 else ' '
        print(
            f'{a + " vs " + b:<62}{diff:>+7.1f} '
            f'{f"[{lo:+.1f}, {hi:+.1f}]":>17}{pval:>8.4f}{star}'
        )
    print('\n* p < 0.05 (exact McNemar, paired). diff is percentage points.')


if __name__ == '__main__':
    main()