"""Train the published FlowMo configurations without overwriting supplied weights."""
from pathlib import Path
import argparse
import json
import os
import subprocess
import sys
import time

CORE = ['flowmo', 'history32', 'history8', 'flowmo_additive', 'flowmo_film']
INITIAL = ['flowmo', 'leworldmodel', 'planet', 'tdmpc2']


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--package', type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--groups', nargs='+', choices=['initial', 'core', 'ood', 'scaling', 'all'], default=['core'])
    parser.add_argument('--methods', nargs='+', choices=sorted(set(CORE + INITIAL)))
    parser.add_argument('--seeds', nargs='+', type=int, default=list(range(1, 9)))
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--smoke', action='store_true', help='One optimizer step; never formal paper statistics')
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()
    root, out = args.package.resolve(), args.output.resolve()
    if out == root or out in root.parents or any(out == root / d or root / d in out.parents for d in ['data', 'experiments', 'paper', 'assets', 'provenance']):
        raise ValueError('Choose a separate training output directory')
    if any(s not in range(1, 9) for s in args.seeds):
        raise ValueError('Formal seeds are 1 through 8')
    cfg = json.loads((root / 'experiments/shared/config/rebuttal_v1.json').read_text())
    groups = ['initial', 'core', 'ood', 'scaling'] if 'all' in args.groups else args.groups
    jobs = []
    for group in groups:
        methods = INITIAL if group == 'initial' else (['flowmo'] if group == 'scaling' else CORE)
        if args.methods:
            methods = [m for m in methods if m in args.methods]
        for method in methods:
            for seed in ([19] if group == 'initial' else args.seeds):
                for count in ([300, 600, 1200] if group == 'scaling' else [2400]):
                    checkpoint = ('paper.pt' if group == 'initial' else
                                  f'rebuttal_scale_e{count:04d}_seed_{seed:04d}.pt' if group == 'scaling' else
                                  f'rebuttal_{group}_seed_{seed:04d}.pt')
                    train_windows = cfg['training']['train_windows'] * count // 2400
                    job = dict(group=group, method=method, seed=seed, episodes=count, checkpoint=checkpoint)
                    cmd = [sys.executable, '-m', 'experiments.train_image_world_models', '--methods', method,
                           '--train-source', str(root / cfg['data']['train_source']), '--test-source', str(root / cfg['data']['test_source']),
                           '--train-episodes', str(64 if args.smoke else count), '--test-episodes', str(64 if args.smoke else 480),
                           '--history-len', '32', '--horizon', '60', '--train-windows', str(512 if args.smoke else train_windows),
                           '--test-windows', str(64 if args.smoke else 24576), '--batch-size', '256',
                           '--microbatch-size', '64', '--eval-batch-size', '64',
                           '--steps', '1' if args.smoke else '20000', '--lr', '0.0003', '--log-every', '1' if args.smoke else '200',
                           '--seed', str(seed), '--seed-mode', 'method_offset' if group == 'initial' else 'paired',
                           '--device', args.device, '--num-workers', '4', '--image-size', '160', '--visual-scale', '2.5',
                           '--render-mode', 'device', '--precision', 'bf16', '--target-mode', 'absolute_normalized',
                           '--heading-weight', '2.0', '--current-pose-weight', '1.0', '--motion-weight', '0.5',
                           '--checkpoint-name', checkpoint, '--checkpoint-interval', '0']
                    if group == 'ood':
                        families = [f for f in cfg['flow_families'] if f not in cfg['heldout_flow_families']]
                        cmd += ['--train-flow-families', *families, '--test-flow-families', *cfg['heldout_flow_families']]
                    if group == 'initial':
                        for flag in ['--microbatch-size', '--eval-batch-size', '--seed-mode']:
                            index = cmd.index(flag)
                            del cmd[index:index + 2]
                    job['source_tree'] = 'protocols/initial/src' if group == 'initial' else '.'
                    jobs.append(dict(job, command=cmd))
    if not jobs:
        raise ValueError('No configurations selected')
    if args.dry_run:
        print(json.dumps({'training_jobs': len(jobs), 'smoke': args.smoke, 'jobs': jobs}, indent=2))
        return
    out.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env['PYTHONPATH'] = str(root) + os.pathsep + env.get('PYTHONPATH', '')
    env['OMP_NUM_THREADS'] = '1'
    env['OPENBLAS_NUM_THREADS'] = '1'
    env['PYTHONPYCACHEPREFIX'] = str(out / 'pycache')
    (out / 'training_plan.json').write_text(json.dumps({'smoke': args.smoke, 'jobs': jobs}, indent=2))
    records = []
    for job in jobs:
        env['PYTHONPATH'] = str(root / job['source_tree']) + os.pathsep + str(root)
        endpoint = out / 'experiments' / job['method'] / 'checkpoint' / job['checkpoint']
        if endpoint.exists():
            raise FileExistsError(f'Checkpoint already exists: {endpoint}; choose another run directory')
        started = time.monotonic()
        log = out / f'{job["group"]}_{job["method"]}_{job["seed"]}_{job["episodes"]}.log'
        with log.open('w') as f:
            subprocess.run(job['command'], cwd=out, env=env, stdout=f, stderr=subprocess.STDOUT, check=True)
        assert endpoint.is_file()
        records.append({k: v for k, v in job.items() if k != 'command'} | {'elapsed_seconds': time.monotonic() - started})
        (out / 'training_receipt.json').write_text(json.dumps({'smoke': args.smoke, 'completed': len(records), 'planned': len(jobs), 'jobs': records}, indent=2))
        print(json.dumps({'completed': len(records), 'planned': len(jobs), 'configuration': records[-1]}), flush=True)


if __name__ == '__main__':
    main()