#!/usr/bin/env python3
"""Run the untouched MT-3 clean/raw/matched/innovation test panel."""
import argparse
import json
import os
import subprocess
import sys


CONDITIONS = ("clean", "raw", "matched", "innovation")
SEEDS = tuple(range(10, 15))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--full_gate", default="results/kp_innovation_full_gate.json")
    parser.add_argument("--condition", choices=("all",) + CONDITIONS,
                        default="all")
    parser.add_argument("--seed", type=int, choices=SEEDS)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--dry_run", action="store_true")
    args = parser.parse_args()
    with open(args.full_gate) as handle:
        gate = json.load(handle)
    if (gate.get("protocol") != "kp_mixed_traffic_full_v1"
            or gate.get("status") != "passed"):
        raise ValueError("MT-2 did not open MT-3")

    conditions = CONDITIONS if args.condition == "all" else (args.condition,)
    seeds = SEEDS if args.seed is None else (args.seed,)
    os.makedirs("results/kp_innovation_confirmation", exist_ok=True)
    for seed in seeds:
        for condition in conditions:
            command = [
                sys.executable, "experiments/conv_run.py",
                "--mode", "kp" if condition == "clean" else "kp_traffic",
            ]
            if condition != "clean":
                command.extend([
                    "--traffic_rule", condition,
                    "--traffic_seed", str(5000 + seed),
                    "--traffic_ratio", "4",
                    "--traffic_calibration_examples", "64",
                    "--learn_P", "1", "--eta_P", "0.1",
                    "--predictor_warmup_steps", "20",
                    "--predictor_every", "16",
                ])
            command.extend([
                "--device", args.device, "--depth", "20", "--width", "16",
                "--seed", str(seed), "--loader_seed", str(seed),
                "--batch_size", "128", "--epochs", "200",
                "--train_limit", "0", "--val_examples", "0",
                "--split_seed", "2027", "--eval_split", "test",
                "--eval_every", "0", "--augment_train", "1",
                "--lr", "0.1", "--output_lr", "0.1",
                "--lr_schedule", "step", "--lr_milestones", "100,150",
                "--lr_gamma", "0.1", "--warmup_epochs", "0",
                "--momentum", "0.9", "--weight_decay", "1e-4",
                "--normalization", "batchnorm", "--a_scale", "1",
                "--alignment_probe", "32",
                "--out", ("results/kp_innovation_confirmation/"
                          f"seed{seed}_{condition}.json"),
            ])
            print(" ".join(command), flush=True)
            if not args.dry_run:
                subprocess.run(command, check=True)


if __name__ == "__main__":
    main()