#!/usr/bin/env python3
"""Launch 8-LLM anti-collapse training runs with post-training evaluation.

This script launches training for three anti-collapse configurations:
1. baseline - Standard anti-collapse weights
2. stronger - 2x weights for aggressive regularization
3. mild - 0.5x weights for gentle regularization

All use post_gnn_layers=0 (dot product scoring only).
Evaluation runs automatically after each training completes.

Usage:
    # Run all three sequentially
    python scripts/launch_8llm_anticollapse.py --all

    # Run specific config
    python scripts/launch_8llm_anticollapse.py --config baseline
    python scripts/launch_8llm_anticollapse.py --config stronger
    python scripts/launch_8llm_anticollapse.py --config mild

    # Skip evaluation
    python scripts/launch_8llm_anticollapse.py --config baseline --no-eval
"""

import argparse
import subprocess
import sys
from pathlib import Path

CONFIGS = {
    "baseline": "configs/anticollapse_8llm_baseline.yaml",
    "stronger": "configs/anticollapse_8llm_stronger.yaml",
    "mild": "configs/anticollapse_8llm_mild.yaml",
}


def run_training(config_name: str, config_path: str, run_eval: bool = True) -> None:
    """Run training for a single config with optional evaluation."""
    print(f"\n{'='*70}")
    print(f"Starting training: {config_name}")
    print(f"Config: {config_path}")
    print(f"Evaluation: {'enabled' if run_eval else 'disabled'}")
    print(f"{'='*70}\n")

    cmd = [
        sys.executable, "-m", "embergraphrouter.run_exp",
        "--config_file", config_path,
    ]

    if not run_eval:
        cmd.append("--no-eval")
    else:
        cmd.append("--eval-advanced")

    subprocess.run(cmd, check=True)


def main():
    parser = argparse.ArgumentParser(description="Launch 8-LLM anti-collapse training")
    parser.add_argument(
        "--config",
        choices=list(CONFIGS.keys()),
        help="Run specific config",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Run all configs sequentially",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print commands without running",
    )
    parser.add_argument(
        "--no-eval",
        action="store_true",
        help="Skip post-training evaluation",
    )
    args = parser.parse_args()

    if not args.config and not args.all:
        parser.print_help()
        print("\nError: Specify --config or --all")
        sys.exit(1)

    run_eval = not args.no_eval
    configs_to_run = list(CONFIGS.items()) if args.all else [(args.config, CONFIGS[args.config])]

    for name, path in configs_to_run:
        if args.dry_run:
            eval_flag = "" if run_eval else " --no-eval"
            print(f"Would run: python -m embergraphrouter.run_exp --config_file {path}{eval_flag}")
        else:
            try:
                run_training(name, path, run_eval=run_eval)
            except subprocess.CalledProcessError as e:
                print(f"Training failed for {name}: {e}")
                if args.all:
                    print("Continuing with next config...")
                else:
                    sys.exit(1)

    print("\n" + "="*70)
    print("All training runs complete!")
    print("="*70)


if __name__ == "__main__":
    main()
