#!/usr/bin/env python3
"""Train S1/S2/S3 (DESIGN §6.4) on one GPU.

python3 scripts/train.py --config configs/train.yaml --stage s2 --seed 42 --out checkpoints/s2_seed42
python3 scripts/train.py --config configs/train.yaml --stage s2 --subset-frac 0.1 --out checkpoints/s2_scan_a --set lambda_rps=1.0
"""

from __future__ import annotations

import argparse
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))

from jev_judge.train_loop import TrainConfig, Trainer  # noqa: E402


def parse_set(items: list[str]) -> dict:
    out = {}
    for it in items or []:
        k, v = it.split("=", 1)
        try:
            import json

            out[k] = json.loads(v)
        except Exception:
            out[k] = v
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/train.yaml")
    ap.add_argument("--stage", choices=["s1", "s2", "s3"])
    ap.add_argument("--seed", type=int)
    ap.add_argument("--out", dest="out_dir")
    ap.add_argument("--model", dest="model_path")
    ap.add_argument("--subset-frac", dest="subset_frac", type=float)
    ap.add_argument("--epochs", type=int)
    ap.add_argument("--max-steps", dest="max_steps", type=int)
    ap.add_argument("--init-from", dest="init_from")
    ap.add_argument("--set", nargs="*", help="extra overrides key=value (JSON values)")
    args = ap.parse_args()
    overrides = {k: v for k, v in vars(args).items() if k not in ("config", "set")}
    overrides.update(parse_set(args.set))
    cfg = TrainConfig.from_yaml(args.config, overrides)
    Trainer(cfg).fit()


if __name__ == "__main__":
    main()