"""Head-to-head comparison of every planner through one harness.

The cross-architecture numbers already in the repo are not comparable to each
other. ``scripts/eval_controller.py`` runs the plain ``WorldModelPolicy``,
which -- per ``eval_planner.py``'s header -- plans from one repeated frame and
zero past action blocks, because ``stable_worldmodel`` declares ``history_len``
and ``action_history`` and implements neither. ``eval_planner.py`` supplies both
via ``HistoryPolicy``. A controller scored by the first and a planner scored by
the second differ in their observations, not just their architecture.

So every row here goes through ``eval_planner.py``, which can evaluate the
controller, the recursive planner, the cross-attention planner, CEM and the
random floor, with history on for all of them and one shared seed -- meaning
identical held-out start/goal pairs and paired McNemar throughout.
"""

import argparse
import subprocess
import sys
from pathlib import Path

PY = sys.executable
ROOT = Path(__file__).resolve().parents[1]
EVAL = 'lejepa_control_2/scripts/eval_planner.py'


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--controller', default='data/runs/ah_hold0.5/controller.pt')
    p.add_argument('--xa', default='data/runs/planner_xa_D/planner.pt')
    p.add_argument('--recursive', default='data/runs/planner_D10_extended/planner.pt')
    p.add_argument('--horizon', type=int, default=5)
    p.add_argument('--receding-horizon', type=int, nargs='+', default=[1, 5])
    p.add_argument('--controller-k', type=int, nargs='+', default=[3, 5])
    p.add_argument('--xa-k', type=int, nargs='+', default=[2, 3, 5])
    p.add_argument('--seeds', type=int, nargs='+', default=[42, 43, 44])
    p.add_argument('--num-eval', type=int, default=50)
    p.add_argument('--goal-offset', type=int, default=25)
    p.add_argument('--out', default='data/runs/eval_compare')
    p.add_argument('--skip-random', action='store_true')
    return p.parse_args()


def main():
    args = parse_args()
    jobs = []

    for rh in args.receding_horizon:
        for seed in args.seeds:
            common = [
                '--horizon', str(args.horizon),
                '--receding-horizon', str(rh),
                '--goal-offset', str(args.goal_offset),
                '--num-eval', str(args.num_eval),
                '--seed', str(seed),
                '--out', args.out,
            ]
            if not args.skip_random and rh == args.receding_horizon[0]:
                jobs.append(['--planner', 'random', *common])

            for k in args.controller_k:
                jobs.append([
                    '--planner', 'controller',
                    '--controller', args.controller,
                    '--refinements', str(k), '--tag', f'ctrl_K{k}', *common,
                ])

            if Path(ROOT / args.recursive).exists():
                jobs.append([
                    '--planner', 'planner', '--checkpoint', args.recursive,
                    '--tag', 'recursive', *common,
                ])

            if Path(ROOT / args.xa).exists():
                for k in args.xa_k:
                    jobs.append([
                        '--planner', 'planner', '--checkpoint', args.xa,
                        '--cycles', str(k), '--tag', f'xa_K{k}', *common,
                    ])

    print(f'{len(jobs)} evaluations -> {args.out}')
    for i, extra in enumerate(jobs, 1):
        print(f'\n[{i}/{len(jobs)}] {" ".join(extra)}', flush=True)
        code = subprocess.run([PY, EVAL, *extra], cwd=ROOT).returncode
        if code != 0:
            print(f'!! exited {code}, continuing', flush=True)

    print(f'\nappended to {Path(args.out) / "results.jsonl"}')


if __name__ == '__main__':
    main()