#!/usr/bin/env python3
"""Launch 2-LLM diagnostic training runs.

This script runs a simplified 2-LLM routing task to test if the GNN
architecture can learn discriminative representations at all.

The 2 LLMs:
- Claude Opus 4.5: Expensive, high quality
- Gemini 2.5 Flash Lite: Cheap, fast

If the GNN fails on this simple binary task, there's a fundamental issue.
If it succeeds, the 8-LLM collapse is due to complexity, not architecture.

Usage:
    # Run both experiments
    python scripts/launch_2llm_diagnostic.py --all

    # Run specific config
    python scripts/launch_2llm_diagnostic.py --config baseline
    python scripts/launch_2llm_diagnostic.py --config no_anticollapse
"""

import argparse
import subprocess
import sys
from pathlib import Path

CONFIGS = {
    "baseline": "configs/gnn_2llm_baseline.yaml",
    "no_anticollapse": "configs/gnn_2llm_no_anticollapse.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"2-LLM DIAGNOSTIC: {config_name}")
    print(f"{'='*70}")
    print(f"Config: {config_path}")
    print(f"Evaluation: {'enabled' if run_eval else 'disabled'}")
    print()

    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 2-LLM diagnostic 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])]

    print("=" * 70)
    print("2-LLM DIAGNOSTIC EXPERIMENT")
    print("=" * 70)
    print("Testing if GNN can learn on simplified binary routing task:")
    print("  - Claude Opus 4.5 (expensive, high quality)")
    print("  - Gemini 2.5 Flash Lite (cheap, fast)")
    print()
    print(f"Configs to run: {[c[0] for c in configs_to_run]}")
    print("=" * 70)

    results = {}
    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)
                results[name] = "SUCCESS"
            except subprocess.CalledProcessError as e:
                print(f"Training failed for {name}: {e}")
                results[name] = "FAILED"
                if args.all:
                    print("Continuing with next config...")
                else:
                    sys.exit(1)

    print("\n" + "=" * 70)
    print("2-LLM DIAGNOSTIC COMPLETE")
    print("=" * 70)
    for name, status in results.items():
        print(f"  {name}: {status}")
    print()
    print("Check model_path/*2llm* for trained weights")
    print("Check *_eval_results.json for evaluation metrics")
    print("=" * 70)


if __name__ == "__main__":
    main()
