"""Summarize the refinement-count sweep, with paired tests.

Success rates at n=50 have a binomial standard error near 7 points, so an
8-point gap between two K values is about one SE and means nothing on its own.
Every row in a sweep shares its held-out start/goal pairs at a fixed seed, so
the comparisons are paired and McNemar's exact test applies -- which is the
standard ``run_report.md`` already holds the rest of the project to.

Reports success vs K per receding horizon, pooled across seeds, and tests each
K against the best-performing K.
"""

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


def mcnemar_exact(a, b):
    """Two-sided exact McNemar on paired boolean outcome lists.

    Returns ``(delta_points, n_discordant, p)``. ``delta`` is ``b`` minus ``a``
    in percentage points.
    """
    assert len(a) == len(b), 'unpaired inputs'
    b_only = sum(1 for x, y in zip(a, b) if not x and y)
    a_only = sum(1 for x, y in zip(a, b) if x and not y)
    n = b_only + a_only
    delta = 100.0 * (b_only - a_only) / len(a)
    if n == 0:
        return delta, 0, 1.0
    k = min(b_only, a_only)
    tail = sum(comb(n, i) for i in range(k + 1)) / (2 ** n)
    return delta, n, min(1.0, 2 * tail)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--results', default='data/runs/eval_refine_sweep/results.jsonl')
    args = p.parse_args()

    rows = [json.loads(line) for line in Path(args.results).read_text().splitlines()]
    rows = [r for r in rows if r.get('refinements') is not None]

    # (rh, K) -> {seed: episode_successes}
    grouped = defaultdict(dict)
    for r in rows:
        grouped[(r['receding_horizon'], r['refinements'])][r['seed']] = (
            r['episode_successes']
        )

    horizons = sorted({rh for rh, _ in grouped})
    for rh in horizons:
        ks = sorted({k for h, k in grouped if h == rh})
        seeds = sorted(set.intersection(
            *[set(grouped[(rh, k)]) for k in ks]
        )) if ks else []
        if not seeds:
            continue

        def pooled(k):
            out = []
            for s in seeds:
                out.extend(grouped[(rh, k)][s])
            return out

        series = {k: pooled(k) for k in ks}
        rates = {k: 100 * sum(v) / len(v) for k, v in series.items()}
        best = max(rates, key=rates.get)

        n = len(series[best])
        print(f'\nreceding_horizon={rh}   seeds={seeds}   n={n} episodes pooled')
        print(f'{"K":>4}{"success%":>10}{"per-seed":>16}'
              f'{"vs best (K=%d)" % best:>28}')
        for k in ks:
            per_seed = ' '.join(
                f'{100 * sum(grouped[(rh, k)][s]) / len(grouped[(rh, k)][s]):.0f}'
                for s in seeds
            )
            if k == best:
                verdict = '(best)'
            else:
                delta, nd, pv = mcnemar_exact(series[k], series[best])
                verdict = f'{delta:+.1f}pt  n_d={nd}  p={pv:.3f}'
            print(f'{k:>4}{rates[k]:>10.1f}{per_seed:>16}{verdict:>28}')

        # the question the sweep exists to answer
        tail = [rates[k] for k in ks if k >= best]
        monotone = all(
            rates[ks[i]] <= rates[ks[i + 1]] + 4 for i in range(len(ks) - 1)
        )
        print(f'  peak at K={best}; '
              f'{"non-decreasing in K (within 4pt noise)" if monotone else "NOT monotone in K"}'
              f'; drop from peak to K={ks[-1]}: {rates[best] - rates[ks[-1]]:+.1f}pt')


if __name__ == '__main__':
    main()