"""Run the baselines and ablations listed in section 11 of the design doc.

Three of them change the objective, so they need their own training runs
(terminal-only, no-support). The rest only change how the trained controller
is executed, so they reuse the main checkpoint. Every sim evaluation uses the
same seed, so all rows share held-out initial-state/goal pairs.
"""

import argparse
import subprocess
import sys
from pathlib import Path

PY = sys.executable
ROOT = Path(__file__).resolve().parents[1]

# objective ablations: each needs its own controller
VARIANTS = {
    'main': [],  # alpha=0.05, lambda_support=0.01
    'terminal_only': ['--alpha', '0'],
    'no_support': ['--lambda-support', '0'],
}


def run(cmd):
    print(f'\n$ {" ".join(str(c) for c in cmd)}', flush=True)
    # one bad row should not discard the hours of runs queued behind it
    code = subprocess.run(cmd, cwd=ROOT).returncode
    if code != 0:
        print(f'!! exited {code}, continuing', flush=True)
    return code == 0


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--steps', type=int, default=20000)
    p.add_argument('--batch-size', type=int, default=128)
    p.add_argument('--lr', type=float, default=3e-4)
    p.add_argument('--num-eval', type=int, default=50)
    p.add_argument('--eval-budget', type=int, default=50)
    p.add_argument('--out', default='data/runs')
    p.add_argument('--skip-train', action='store_true')
    p.add_argument('--skip-cem', action='store_true')
    return p.parse_args()


def main():
    args = parse_args()
    out = Path(args.out)

    def evaluate(extra, tag_dir='eval'):
        run(
            [PY, 'scripts/eval_controller.py',
             '--num-eval', str(args.num_eval),
             '--eval-budget', str(args.eval_budget),
             '--out', str(out / tag_dir), *extra]
        )

    # Everything that reuses the main checkpoint runs first: those are the
    # headline numbers and cost minutes, while each objective ablation costs
    # its own multi-hour training run.

    # --- baseline: LeWM + CEM ---------------------------------------------
    if not args.skip_cem:
        evaluate(['--planner', 'cem'])

    main_ckpt = ['--controller', str(out / 'controller' / 'controller.pt')]

    # --- refinement count: K=0 (one-pass) through K=5 ---------------------
    for k in (0, 1, 2, 3, 5):
        evaluate([*main_ckpt, '--refinements', str(k)])

    # --- execute one block vs the whole plan ------------------------------
    evaluate([*main_ckpt, '--receding-horizon', '5'])

    # --- objective ablations need their own training runs -----------------
    for name, extra in VARIANTS.items():
        if name == 'main':
            continue
        ckpt_dir = out / f'abl_{name}'
        if not args.skip_train:
            run(
                [PY, 'scripts/train_controller.py',
                 '--steps', str(args.steps),
                 '--batch-size', str(args.batch_size),
                 '--lr', str(args.lr),
                 '--log-every', '500', '--val-every', '5000',
                 '--out', str(ckpt_dir), *extra]
            )
        ckpt = ckpt_dir / 'controller.pt'
        if ckpt.exists():
            evaluate(['--controller', str(ckpt)])

    print(f'\nall results appended to {out / "eval" / "results.jsonl"}')


if __name__ == '__main__':
    main()