#!/usr/bin/env python3
"""Launch loss function comparison sweep on multiple GPUs.

Compares 5 loss variants with identical architecture/hyperparameters:
1. listwise - Pure listwise KL (control, 81.02% baseline)
2. bradley_terry - Pure Bradley-Terry pairwise
3. lambda_rank - Pure LambdaRank
4. hybrid_bradley_terry - Listwise + Bradley-Terry
5. hybrid_lambda_rank - Listwise + LambdaRank

Usage:
    python scripts/launch_loss_sweep.py --gpus 0,1,2,3,4
"""

from __future__ import annotations

import argparse
import logging
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

LOSS_VARIANTS = [
    {"name": "listwise", "config": "configs/loss_sweep_listwise.yaml"},
    {"name": "bradley_terry", "config": "configs/loss_sweep_bradley_terry.yaml"},
    {"name": "lambda_rank", "config": "configs/loss_sweep_lambda_rank.yaml"},
    {"name": "hybrid_bradley_terry", "config": "configs/loss_sweep_hybrid_bradley_terry.yaml"},
    {"name": "hybrid_lambda_rank", "config": "configs/loss_sweep_hybrid_lambda_rank.yaml"},
]


def launch_variant(variant: dict, gpu_id: int, run_dir: Path) -> subprocess.Popen:
    """Launch a single loss variant on a specific GPU."""
    name = variant["name"]
    config = variant["config"]
    log_file = run_dir / f"{name}.log"

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    cmd = [
        sys.executable,
        "run_exp.py",
        "--config", config,
    ]

    with open(log_file, "w") as f:
        proc = subprocess.Popen(
            cmd,
            stdout=f,
            stderr=subprocess.STDOUT,
            env=env,
        )

    return proc


def main():
    parser = argparse.ArgumentParser(description="Launch loss function sweep")
    parser.add_argument(
        "--gpus",
        type=str,
        default="0,1,2,3,4",
        help="Comma-separated GPU IDs to use (default: 0,1,2,3,4)",
    )
    args = parser.parse_args()

    gpu_ids = [int(g) for g in args.gpus.split(",")]

    if len(gpu_ids) < len(LOSS_VARIANTS):
        logger.warning(
            f"Only {len(gpu_ids)} GPUs available for {len(LOSS_VARIANTS)} variants. "
            "Some variants will be skipped."
        )

    # Create run directory
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = Path(f"runs/loss_sweep_{timestamp}")
    run_dir.mkdir(parents=True, exist_ok=True)

    logger.info("=" * 60)
    logger.info("LOSS FUNCTION SWEEP")
    logger.info("=" * 60)
    logger.info(f"GPUs: {gpu_ids}")
    logger.info(f"Variants: {[v['name'] for v in LOSS_VARIANTS[:len(gpu_ids)]]}")

    # Launch variants
    processes = []
    for i, (variant, gpu_id) in enumerate(zip(LOSS_VARIANTS, gpu_ids)):
        logger.info(f"Launching {variant['name']} on GPU {gpu_id}")
        proc = launch_variant(variant, gpu_id, run_dir)
        processes.append((variant["name"], proc))
        logger.info(f"  Started PID {proc.pid}, log: {run_dir / variant['name']}.log")

    logger.info(f"\nMonitoring... (logs in {run_dir})")

    # Wait for all processes
    for name, proc in processes:
        proc.wait()
        if proc.returncode == 0:
            logger.info(f"{name} completed successfully")
        else:
            logger.error(f"{name} failed with code {proc.returncode}")

    logger.info("=" * 60)
    logger.info("SWEEP COMPLETE")
    logger.info("=" * 60)

    # Print summary
    logger.info("\nResults summary (check individual logs for metrics):")
    for name, _ in processes:
        log_file = run_dir / f"{name}.log"
        # Extract best frac_oracle from log
        try:
            with open(log_file) as f:
                content = f.read()
                # Look for best score
                import re
                matches = re.findall(r"New best model saved \(score=([0-9.]+)\)", content)
                if matches:
                    best = max(float(m) for m in matches)
                    logger.info(f"  {name}: best_score={best:.4f}")
                else:
                    logger.info(f"  {name}: no best model saved")
        except Exception as e:
            logger.info(f"  {name}: error reading log - {e}")


if __name__ == "__main__":
    main()
