#!/usr/bin/env python3
"""Launch multi-embedding sweep on B200 GPUs.

Experiments:
- Gemini text-embedding-004: 256, 512, 768 dims
- OpenAI text-embedding-3-large: 256, 512, 768, 1536, 3072 dims

Architecture adapts hidden dimension based on input embedding size (no padding):
- 256-dim  -> 256 hidden
- 512-dim  -> 384 hidden
- 768-dim  -> 512 hidden
- 1536-dim -> 768 hidden
- 3072-dim -> 1024 hidden

Prerequisites:
- configs/sweep_*.yaml configs created
- configs/llm_embedding_*.pkl files generated
- data/router_data_multi_embed.parquet with all embedding columns

Usage:
    python scripts/launch_multi_embed_sweep.py [--dry-run] [--gpus 0,1,2,3,4,5,6,7]
"""

import argparse
import logging
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent

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


EMBEDDING_VARIANTS = [
    {"name": "gemini_256", "dim": 256, "hidden": 256},
    {"name": "gemini_512", "dim": 512, "hidden": 384},
    {"name": "gemini_768", "dim": 768, "hidden": 512},
    {"name": "openai_large_256", "dim": 256, "hidden": 256},
    {"name": "openai_large_512", "dim": 512, "hidden": 384},
    {"name": "openai_large_768", "dim": 768, "hidden": 512},
    {"name": "openai_large_1536", "dim": 1536, "hidden": 768},
    {"name": "openai_large_3072", "dim": 3072, "hidden": 1024},
]


def get_available_gpus() -> List[int]:
    """Detect available CUDA GPUs."""
    try:
        import torch
        if torch.cuda.is_available():
            return list(range(torch.cuda.device_count()))
    except ImportError:
        pass
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
            capture_output=True, text=True, check=True,
        )
        return [int(idx.strip()) for idx in result.stdout.strip().split("\n") if idx.strip()]
    except (subprocess.CalledProcessError, FileNotFoundError):
        return []


def verify_prerequisites(variants: List[Dict]) -> bool:
    """Verify all required files exist."""
    all_ok = True

    # Check parquet
    parquet = _MODULE_DIR / "data" / "router_data_multi_embed.parquet"
    if not parquet.exists():
        logger.error(f"Missing: {parquet}")
        all_ok = False

    for v in variants:
        # Check config
        config = _MODULE_DIR / "configs" / f"sweep_{v['name']}.yaml"
        if not config.exists():
            logger.error(f"Missing config: {config}")
            all_ok = False

        # Check LLM embeddings
        llm_emb = _MODULE_DIR / "configs" / f"llm_embedding_{v['name']}.pkl"
        if not llm_emb.exists():
            logger.error(f"Missing LLM embeddings: {llm_emb}")
            all_ok = False

    return all_ok


def launch_training(
    variant: Dict,
    gpu_id: int,
    log_dir: Path,
    dry_run: bool = False,
) -> Optional[subprocess.Popen]:
    """Launch training on a specific GPU."""
    name = variant["name"]
    config_path = _MODULE_DIR / "configs" / f"sweep_{name}.yaml"
    log_file = log_dir / f"{name}.log"

    cmd = [sys.executable, str(_MODULE_DIR / "run_exp.py"), "--config_file", str(config_path)]
    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    logger.info(f"Launching {name} on GPU {gpu_id} (dim={variant['dim']}, hidden={variant['hidden']})")

    if dry_run:
        logger.info(f"  [DRY] CUDA_VISIBLE_DEVICES={gpu_id} python run_exp.py --config_file {config_path}")
        return None

    with open(log_file, "w") as f:
        proc = subprocess.Popen(cmd, env=env, stdout=f, stderr=subprocess.STDOUT, cwd=str(_MODULE_DIR))
    logger.info(f"  Started PID {proc.pid}, log: {log_file}")
    return proc


def monitor_processes(processes: Dict[str, subprocess.Popen], log_dir: Path, poll_interval: int = 60) -> Dict[str, int]:
    """Monitor running processes."""
    results = {}
    running = set(processes.keys())

    while running:
        time.sleep(poll_interval)
        for name in list(running):
            ret = processes[name].poll()
            if ret is not None:
                running.remove(name)
                results[name] = ret
                status = "SUCCESS" if ret == 0 else f"FAILED({ret})"
                logger.info(f"[{status}] {name}")
        if running:
            logger.info(f"Running ({len(running)}): {', '.join(sorted(running))}")
    return results


def main():
    parser = argparse.ArgumentParser(description="Launch multi-embedding sweep")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--gpus", type=str, default=None, help="Comma-separated GPU IDs")
    parser.add_argument("--variants", type=str, default=None, help="Comma-separated variant names")
    parser.add_argument("--log-dir", type=str, default=None)
    args = parser.parse_args()

    logger.info("=" * 60)
    logger.info("MULTI-EMBEDDING SWEEP")
    logger.info("=" * 60)

    # Filter variants
    variants = EMBEDDING_VARIANTS
    if args.variants:
        names = [v.strip() for v in args.variants.split(",")]
        variants = [v for v in EMBEDDING_VARIANTS if v["name"] in names]

    # Get GPUs
    gpus = [int(g) for g in args.gpus.split(",")] if args.gpus else get_available_gpus()
    if not gpus:
        logger.error("No GPUs available")
        sys.exit(1)

    logger.info(f"GPUs: {gpus}")
    logger.info(f"Variants: {[v['name'] for v in variants]}")

    # Verify prerequisites
    if not verify_prerequisites(variants):
        sys.exit(1)

    # Create log dir
    if args.log_dir:
        log_dir = Path(args.log_dir)
    else:
        log_dir = _MODULE_DIR / "runs" / f"embed_sweep_{time.strftime('%Y%m%d_%H%M%S')}"
    log_dir.mkdir(parents=True, exist_ok=True)

    # Launch
    processes = {}
    for i, variant in enumerate(variants):
        gpu_id = gpus[i % len(gpus)]
        proc = launch_training(variant, gpu_id, log_dir, args.dry_run)
        if proc:
            processes[variant["name"]] = proc

    if args.dry_run:
        return

    # Monitor
    logger.info(f"\nMonitoring... (logs in {log_dir})")
    try:
        results = monitor_processes(processes, log_dir)
        logger.info("\n" + "=" * 60)
        logger.info("SWEEP COMPLETE")
        for name, ret in sorted(results.items()):
            logger.info(f"  {name}: {'OK' if ret == 0 else f'FAIL({ret})'}")
    except KeyboardInterrupt:
        logger.info("\nStopped monitoring. Training continues in background.")


if __name__ == "__main__":
    main()
