#!/usr/bin/env python3
"""Prepare optional multi-seed rerun commands for GRL sensitivity checks."""

from __future__ import annotations

import argparse
import subprocess
from pathlib import Path


def phase_command(seed: int, phase_balanced: bool) -> list[str]:
    if phase_balanced:
        return [
            "python",
            "scripts/snr_transfer_phase_balanced_experiment.py",
            "--out-dir",
            f"outputs/snr_transfer_phase_balanced_seed{seed}",
            "--cache-dir",
            "outputs/snr_transfer_phase_balanced_cache",
            "--seed",
            str(seed),
            "--init-modes",
            "finetune",
            "scratch",
            "--scratch-train-steps",
            "10000",
            "--resume",
        ]
    return [
        "python",
        "scripts/snr_transfer_experiment.py",
        "--out-dir",
        f"outputs/snr_transfer_seed{seed}",
        "--seed",
        str(seed),
        "--resume",
    ]


def dispersion_command(seed: int) -> list[str]:
    return [
        "python",
        "scripts/disp_snr_transfer_experiment.py",
        "--out-dir",
        f"outputs/disp_snr_transfer_seed{seed}",
        "--seed",
        str(seed),
        "--epochs",
        "5",
        "--batch-size",
        "256",
        "--device",
        "cpu",
    ]


def run_or_print(cmd: list[str], dry_run: bool) -> None:
    print(" ".join(cmd))
    if not dry_run:
        subprocess.run(cmd, check=True)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seeds", nargs="+", type=int, default=[20260609, 20260610, 20260611])
    parser.add_argument("--task", choices=["phase", "dispersion", "both"], default="both")
    parser.add_argument("--dry-run", action="store_true", default=False)
    parser.add_argument("--skip-existing", action="store_true", default=False)
    parser.add_argument(
        "--phase-balanced",
        action="store_true",
        help="Run the phase-aware P/S-balanced SNR experiment instead of the original record-level SNR experiment.",
    )
    args = parser.parse_args()
    for seed in args.seeds:
        if args.task in ("phase", "both"):
            phase_prefix = "snr_transfer_phase_balanced_seed" if args.phase_balanced else "snr_transfer_seed"
            out = Path(f"outputs/{phase_prefix}{seed}/summary.json")
            if args.skip_existing and out.exists():
                print(f"# skip existing {out}")
            else:
                run_or_print(phase_command(seed, args.phase_balanced), args.dry_run)
        if args.task in ("dispersion", "both"):
            out = Path(f"outputs/disp_snr_transfer_seed{seed}/summary.json")
            if args.skip_existing and out.exists():
                print(f"# skip existing {out}")
            else:
                run_or_print(dispersion_command(seed), args.dry_run)


if __name__ == "__main__":
    main()