#!/usr/bin/env python3
"""Launch v15 multi-seed GNN training with post-training evaluation.

v15 introduces a learned score head (Linear(512, 1)) to replace the fixed
mean-pooling that caused logit collapse in v14. Also enables focal loss,
needle weighting, and hard negative mining.

Usage:
    # Full 3-seed run (2026, 42, 1337) with evaluation
    python scripts/launch_v15_training.py

    # Single seed for smoke testing
    python scripts/launch_v15_training.py --seeds 42

    # Custom epochs for quick validation
    python scripts/launch_v15_training.py --seeds 42 --epochs 50

    # Dry run (print commands only)
    python scripts/launch_v15_training.py --dry-run
"""

import argparse
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path

_SCRIPT_DIR = Path(__file__).resolve().parent
_PACKAGE_DIR = _SCRIPT_DIR.parent
_CONFIG_PATH = _PACKAGE_DIR / "configs" / "config_v15_2026.yaml"

# v14 baselines for comparison (quality-only, perf=1.0)
V14_BASELINES = {
    "best_single_llm": 0.846,  # GPT-5.2 Codex
    "v14_seed_2026": 0.859,
    "v14_seed_42": 0.860,
    "v14_seed_1337": 0.853,
    "v14_mean": 0.857,
    "v13_mean": 0.754,
    "random": 0.764,
}

DEFAULT_SEEDS = [2026, 42, 1337]


def run_training(
    seed: int,
    checkpoint_dir: Path,
    epochs: int | None = None,
    dry_run: bool = False,
    run_eval: bool = True,
) -> None:
    """Run training for a single seed."""
    seed_dir = checkpoint_dir / f"seed_{seed}"
    seed_dir.mkdir(parents=True, exist_ok=True)

    model_path = seed_dir / "model.pth"

    cmd = [
        sys.executable, "-m", "embergraphrouter.run_exp",
        "--config_file", str(_CONFIG_PATH),
    ]

    # Override seed and model_path via config patching
    overrides = f"--seed {seed} --model_path {model_path}"
    if epochs is not None:
        overrides += f" --train_epoch {epochs}"

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

    if dry_run:
        print(f"  [seed={seed}] Would run: {' '.join(cmd)}")
        print(f"  Model output: {model_path}")
        return

    print(f"\n{'='*70}")
    print(f"Training seed={seed} | epochs={epochs or 300} | output={seed_dir}")
    print(f"{'='*70}\n")

    # Build a temporary config with seed/model_path overrides
    import yaml
    with open(_CONFIG_PATH) as f:
        config = yaml.safe_load(f)

    config["seed"] = seed
    config["model_path"] = str(model_path)
    if epochs is not None:
        config["train_epoch"] = epochs

    tmp_config = seed_dir / "config.yaml"
    with open(tmp_config, "w") as f:
        yaml.dump(config, f, default_flow_style=False)

    run_cmd = [
        sys.executable, "-m", "embergraphrouter.run_exp",
        "--config_file", str(tmp_config),
    ]
    if run_eval:
        run_cmd.append("--eval-advanced")
    else:
        run_cmd.append("--no-eval")

    subprocess.run(run_cmd, check=True, cwd=str(_PACKAGE_DIR))


def print_comparison(checkpoint_dir: Path, seeds: list[int]) -> None:
    """Print v15 results compared to v14 baselines."""
    print(f"\n{'='*70}")
    print("v15 Training Results")
    print(f"{'='*70}")
    print(f"Checkpoint dir: {checkpoint_dir}")
    print()

    print("v14 Baselines:")
    for name, val in V14_BASELINES.items():
        print(f"  {name:25s} {val:.3f}")
    print()

    print("v15 Results:")
    for seed in seeds:
        seed_dir = checkpoint_dir / f"seed_{seed}"
        model_path = seed_dir / "model.pth"
        if model_path.exists():
            size_mb = model_path.stat().st_size / (1024 * 1024)
            print(f"  seed_{seed:>4d}: model saved ({size_mb:.1f} MB)")

            # Check for score_head in state_dict
            try:
                import torch
                state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
                has_score_head = "score_head.weight" in state_dict
                print(f"    score_head: {has_score_head}")
                if has_score_head:
                    weight = state_dict["score_head.weight"]
                    print(f"    score_head weight norm: {weight.norm().item():.4f}")
                    print(f"    score_head weight std:  {weight.std().item():.6f}")
            except Exception:
                pass
        else:
            print(f"  seed_{seed:>4d}: NOT FOUND")
    print()
    print("Run eval_difficulty_aware.py on each seed for frac_oracle scores.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Launch v15 multi-seed GNN training"
    )
    parser.add_argument(
        "--seeds",
        type=int,
        nargs="+",
        default=DEFAULT_SEEDS,
        help="Seeds to train (default: 2026 42 1337)",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=None,
        help="Override train_epoch (default: 300 from config)",
    )
    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",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Override checkpoint output directory",
    )
    args = parser.parse_args()

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    if args.output_dir:
        checkpoint_dir = Path(args.output_dir)
    else:
        checkpoint_dir = _PACKAGE_DIR / "checkpoints" / f"v15_{timestamp}"

    if not args.dry_run:
        checkpoint_dir.mkdir(parents=True, exist_ok=True)
        # Copy config for reproducibility
        shutil.copy2(_CONFIG_PATH, checkpoint_dir / "config_v15_2026.yaml")

    print(f"v15 Training Launch")
    print(f"  Seeds: {args.seeds}")
    print(f"  Epochs: {args.epochs or 300}")
    print(f"  Output: {checkpoint_dir}")
    print(f"  Config: {_CONFIG_PATH}")
    print()

    run_eval = not args.no_eval

    for seed in args.seeds:
        try:
            run_training(
                seed=seed,
                checkpoint_dir=checkpoint_dir,
                epochs=args.epochs,
                dry_run=args.dry_run,
                run_eval=run_eval,
            )
        except subprocess.CalledProcessError as e:
            print(f"\nTraining FAILED for seed={seed}: {e}")
            if len(args.seeds) > 1:
                print("Continuing with next seed...")
            else:
                sys.exit(1)

    if not args.dry_run:
        print_comparison(checkpoint_dir, args.seeds)


if __name__ == "__main__":
    main()
