"""Summarize eval runs into the tables the design doc asks for.

Reads the jsonl written by ``scripts/eval_controller.py``. The central claim
is matching CEM's success at a fraction of the predictor evaluations, so cost
is reported as predictor rows per episode alongside wall-clock: CEM batches
its samples, so the two ratios differ by an order of magnitude and only
stating both is honest.

Rows are compared against the CEM baseline *at the same execution length*.
Comparing against CEM at receding-horizon 1 flatters the controller, because
the fixed-terminal objective makes every planner procrastinate under short
execution.
"""

import argparse
import json
from pathlib import Path


def refinements(row):
    k = row.get('refinements')
    if k is not None:
        return k
    if row['planner'].startswith('controller'):
        return int(row['planner'].split('_K')[-1])
    return None


def variant(row):
    """Which trained checkpoint a row came from."""
    ckpt = (row.get('checkpoint') or '').replace('\\', '/')
    for name in ('terminal_only', 'no_support',
                 'ah_hold0.5', 'ah_hold0.0', 'ah_hold1.0'):
        if name in ckpt:
            return name
    return 'main'


def label(row):
    tag = row['planner']
    if row.get('receding_horizon', 1) != 1:
        tag += f'+exec{row["receding_horizon"]}'
    v = variant(row)
    if v != 'main':
        tag += f'[{v}]'
    return tag


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

    rows = [
        json.loads(x)
        for x in Path(args.results).read_text().splitlines()
        if x.strip()
    ]
    if not rows:
        print('no results yet')
        return

    # baseline per execution length, so cost ratios are never computed across
    # schedules that make different numbers of solver calls
    cem_by_exec = {
        r.get('receding_horizon', 1): r
        for r in rows if r['planner'].startswith('cem')
    }

    head = (
        f'{"planner":<34}{"exec":>5}{"success%":>10}{"first_d":>9}'
        f'{"rows/ep":>10}{"sec/ep":>8}'
    )
    print(head)
    print('-' * len(head))
    for r in rows:
        d = r.get('first_terminal_distance')
        print(
            f'{label(r):<34}{r.get("receding_horizon", 1):>5}'
            f'{r["success_rate"]:>10.1f}'
            f'{(f"{d:.4f}" if d is not None else "-"):>9}'
            f'{r["predictor_rows_per_episode"]:>10.0f}'
            f'{r["seconds_per_episode"]:>8.2f}'
        )

    # --- K x execution grid, per trained variant ---------------------------
    grids = {}
    for r in rows:
        if not r['planner'].startswith('controller'):
            continue
        grids.setdefault(variant(r), {})[
            (refinements(r), r.get('receding_horizon', 1))
        ] = r['success_rate']

    for name, grid in grids.items():
        ks = sorted({k for k, _ in grid})
        ms = sorted({m for _, m in grid})
        print(f'\nsuccess% [{name}] — refinements K (rows) x execution (cols)')
        print(f'{"K":>3}' + ''.join(f'{m:>8}' for m in ms))
        for k in ks:
            cells = ''.join(
                f'{grid[(k, m)]:>8.1f}' if (k, m) in grid else f'{"-":>8}'
                for m in ms
            )
            print(f'{k:>3}{cells}')

    # --- matched-execution comparison --------------------------------------
    print('\nvs CEM at the same execution length:')
    for r in rows:
        if not r['planner'].startswith('controller'):
            continue
        base = cem_by_exec.get(r.get('receding_horizon', 1))
        if base is None:
            continue
        rows_x = base['predictor_rows_per_episode'] / max(
            r['predictor_rows_per_episode'], 1e-9
        )
        time_x = base['seconds_per_episode'] / max(
            r['seconds_per_episode'], 1e-9
        )
        print(
            f'  {label(r):<36}{r["success_rate"]:>6.1f}% vs '
            f'{base["success_rate"]:>5.1f}%   '
            f'{rows_x:>7.0f}x fewer rows   {time_x:>5.1f}x faster'
        )


if __name__ == '__main__':
    main()