#!/usr/bin/env python3
"""Run the three Phase 2.1A v2 Robometer-Low-LoRA seeds serially with resume."""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
TRAINER = ROOT / "scripts" / "train_phase21a_robometer_low_lora.py"
def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--snapshot", type=Path, required=True)
    parser.add_argument(
        "--manifest",
        type=Path,
        default=ROOT / "manifests" / "phase2_remediation_pairs_v1.jsonl",
    )
    parser.add_argument(
        "--processed-root",
        type=Path,
        default=ROOT / "artifacts" / "data_views" / "phase1_official" / "processed_datasets",
    )
    parser.add_argument("--seeds", type=int, nargs="+", default=[7, 17, 27])
    parser.add_argument("--optimizer-steps", type=int, default=3000)
    parser.add_argument("--micro-batch-pairs", type=int, default=32)
    parser.add_argument("--gradient-accumulation", type=int, default=2)
    parser.add_argument("--lora-lr", type=float, default=2.0e-6)
    parser.add_argument("--head-lr", type=float, default=1.0e-4)
    parser.add_argument("--checkpoint-every", type=int, default=250)
    parser.add_argument("--eval-every", type=int, default=250)
    parser.add_argument("--eval-micro-batch-pairs", type=int, default=32)
    parser.add_argument("--early-stop-min-steps", type=int, default=1000)
    parser.add_argument("--early-stop-patience-evals", type=int, default=4)
    parser.add_argument("--early-stop-min-delta", type=float, default=5.0e-3)
    args = parser.parse_args()
    for seed in args.seeds:
        output = args.output_root / f"seed_{seed}"
        if (output / "summary.json").is_file():
            continue
        command = [
            sys.executable,
            str(TRAINER),
            "--snapshot", str(args.snapshot),
            "--manifest", str(args.manifest),
            "--processed-root", str(args.processed_root),
            "--seed", str(seed),
            "--optimizer-steps", str(args.optimizer_steps),
            "--micro-batch-pairs", str(args.micro_batch_pairs),
            "--gradient-accumulation", str(args.gradient_accumulation),
            "--lora-lr", str(args.lora_lr),
            "--head-lr", str(args.head_lr),
            "--checkpoint-every", str(args.checkpoint_every),
            "--eval-every", str(args.eval_every),
            "--eval-micro-batch-pairs", str(args.eval_micro_batch_pairs),
            "--early-stop-min-steps", str(args.early_stop_min_steps),
            "--early-stop-patience-evals", str(args.early_stop_patience_evals),
            "--early-stop-min-delta", str(args.early_stop_min_delta),
            "--output-dir", str(output),
        ]
        if (output / "latest_training_state.pt").is_file():
            command.append("--resume")
        subprocess.run(command, check=True)


if __name__ == "__main__":
    main()