#!/usr/bin/env python3
"""Robust ablation sweep launcher with GPU failure recovery.

This script launches multiple experiments in parallel with:
- Pre-launch GPU health checks
- Automatic retry on alternative GPUs
- Progress monitoring
- Checkpoint resumption support

Usage:
    python scripts/launch_ablation_sweep.py --configs configs/anticollapse_ablation_*.yaml
"""

import argparse
import subprocess
import time
from pathlib import Path
from typing import Optional


def check_gpu_availability(gpu_id: int) -> bool:
    """Test if a GPU is available and healthy."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "-i", str(gpu_id), "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            return False
        memory_used = int(result.stdout.strip())
        # Consider GPU available if using less than 1GB
        return memory_used < 1024
    except Exception:
        return False


def get_available_gpus(exclude: Optional[list[int]] = None) -> list[int]:
    """Get list of available GPU IDs, excluding specified ones."""
    exclude = exclude or []
    available = []
    for gpu_id in range(8):  # Assume 8 GPUs max
        if gpu_id not in exclude and check_gpu_availability(gpu_id):
            available.append(gpu_id)
    return available


def launch_experiment(
    config_path: Path,
    gpu_id: int,
    log_path: Path,
    python_cmd: str = ".venv/bin/python3",
) -> Optional[subprocess.Popen]:
    """Launch a single experiment on specified GPU."""
    cmd = f"CUDA_VISIBLE_DEVICES={gpu_id} nohup {python_cmd} run_exp.py --config_file {config_path} > {log_path} 2>&1 &"

    try:
        # Use shell=True for nohup background execution
        subprocess.run(cmd, shell=True, check=True)
        # Give it a moment to start
        time.sleep(2)

        # Verify it started by checking if log file has content
        if log_path.exists() and log_path.stat().st_size > 0:
            return True
        return False
    except subprocess.CalledProcessError:
        return False


def check_experiment_alive(log_path: Path, timeout_minutes: int = 5) -> bool:
    """Check if experiment is still running based on log file updates."""
    if not log_path.exists():
        return False

    mtime = log_path.stat().st_mtime
    age_minutes = (time.time() - mtime) / 60
    return age_minutes < timeout_minutes


def main():
    parser = argparse.ArgumentParser(description="Launch ablation sweep with GPU recovery")
    parser.add_argument("--configs", nargs="+", required=True, help="Config files to run")
    parser.add_argument("--exclude-gpus", nargs="*", type=int, default=[1], help="GPUs to exclude (default: 1)")
    parser.add_argument("--log-dir", type=Path, default=Path("logs"), help="Log directory")
    parser.add_argument("--max-retries", type=int, default=3, help="Max retries per experiment")
    args = parser.parse_args()

    args.log_dir.mkdir(exist_ok=True)

    # Get available GPUs
    available_gpus = get_available_gpus(exclude=args.exclude_gpus)
    print(f"Available GPUs (excluding {args.exclude_gpus}): {available_gpus}")

    if len(available_gpus) < len(args.configs):
        print(f"Warning: Only {len(available_gpus)} GPUs for {len(args.configs)} experiments")

    # Launch experiments
    experiments = {}
    for i, config in enumerate(args.configs):
        config_path = Path(config)
        name = config_path.stem
        log_path = args.log_dir / f"{name}.log"

        gpu_id = available_gpus[i % len(available_gpus)]

        print(f"Launching {name} on GPU {gpu_id}...")
        success = launch_experiment(config_path, gpu_id, log_path)

        if success:
            experiments[name] = {"gpu": gpu_id, "log": log_path, "retries": 0}
            print(f"  ✓ Started on GPU {gpu_id}")
        else:
            print(f"  ✗ Failed to start on GPU {gpu_id}")
            experiments[name] = {"gpu": None, "log": log_path, "retries": 1, "pending": True}

    print(f"\nLaunched {len([e for e in experiments.values() if e.get('gpu')])} experiments")
    print("Use 'tail -f logs/*.log' to monitor progress")


if __name__ == "__main__":
    main()
