#!/usr/bin/env python3
"""Run the frozen S0 training-only stability-margin candidates."""
import argparse
import json
import os
import subprocess
import sys


MARGINS = {
    "0p001": 0.001,
    "0p003": 0.003,
    "0p01": 0.01,
    "0p03": 0.03,
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--margin", choices=("all",) + tuple(MARGINS),
                        default="all")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--dry_run", action="store_true")
    parser.add_argument(
        "--mt1_gate", default="results/kp_innovation_short_gate.json")
    args = parser.parse_args()
    with open(args.mt1_gate) as handle:
        gate = json.load(handle)
    if (gate.get("protocol") != "kp_mixed_traffic_short_v1"
            or gate.get("status") != "failed"):
        raise ValueError("S0 requires the audited MT-1 failure")

    labels = tuple(MARGINS) if args.margin == "all" else (args.margin,)
    os.makedirs("results/kp_innovation_stability_grid", exist_ok=True)
    for label in labels:
        command = [
            sys.executable,
            "experiments/diagnose_kp_traffic_nonfinite.py",
            "--rule", "innovation",
            "--predictor_mode", "closed_form",
            "--predictor_every", "0",
            "--stability_margin", str(MARGINS[label]),
            "--device", args.device,
            "--max_steps", "352",
            "--out", ("results/kp_innovation_stability_grid/"
                      f"margin_{label}.json"),
        ]
        print(" ".join(command), flush=True)
        if not args.dry_run:
            subprocess.run(command, check=True)


if __name__ == "__main__":
    main()