#!/usr/bin/env python3
"""Robust Qwen3 sweep launcher with GPU health detection and resumption.

Features:
- Detects healthy GPUs before launching
- Tracks job status in a state file for resumption
- Requeues failed jobs to healthy GPUs
- Captures comprehensive metrics to CSV

Usage:
    python scripts/launch_qwen3_sweep.py --check-gpus     # Check GPU health
    python scripts/launch_qwen3_sweep.py --launch         # Launch sweep
    python scripts/launch_qwen3_sweep.py --status         # Check status
    python scripts/launch_qwen3_sweep.py --resume         # Resume failed jobs
"""

import argparse
import json
import subprocess
import time
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path


@dataclass
class JobState:
    config_name: str
    gpu_id: int
    pid: int | None
    status: str  # pending, running, completed, failed
    start_time: str | None
    end_time: str | None
    log_file: str
    best_frac_oracle: float | None
    best_top1_acc: float | None
    best_val_acc: float | None


SWEEP_CONFIGS = [
    "qwen3_sweep_01_focal_lr1e4",
    "qwen3_sweep_02_focal_lr5e4",
    "qwen3_sweep_03_hybrid_lr1e4",
    "qwen3_sweep_04_hybrid_lr5e4",
    "qwen3_sweep_05_bce_lr1e4",
    "qwen3_sweep_06_bce_lr5e4",
    "qwen3_sweep_07_listwise_lr1e4",
    "qwen3_sweep_08_focal_no_difficulty",
]

STATE_FILE = Path("sweep_state.json")
LOGS_DIR = Path("logs")
CONFIGS_DIR = Path("configs")


def check_gpu_health() -> list[int]:
    """Check which GPUs are healthy (no CUDA errors, memory available)."""
    healthy_gpus = []

    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,memory.used,memory.total,gpu_bus_id",
             "--format=csv,noheader,nounits"],
            capture_output=True,
            text=True,
            timeout=30,
        )

        for line in result.stdout.strip().split("\n"):
            if not line.strip():
                continue
            parts = [p.strip() for p in line.split(",")]
            gpu_id = int(parts[0])
            mem_used = int(parts[1])
            mem_total = int(parts[2])

            # GPU is healthy if <10% memory used (allows for driver overhead)
            if mem_used < mem_total * 0.1:
                healthy_gpus.append(gpu_id)
                print(f"  GPU {gpu_id}: HEALTHY ({mem_used}/{mem_total} MiB used)")
            else:
                print(f"  GPU {gpu_id}: BUSY ({mem_used}/{mem_total} MiB used)")

    except Exception as e:
        print(f"Error checking GPUs: {e}")

    return healthy_gpus


def test_gpu_cuda(gpu_id: int) -> bool:
    """Test if a GPU can run CUDA operations."""
    test_script = f'''
import torch
torch.cuda.set_device({gpu_id})
x = torch.randn(100, 100, device="cuda")
y = x @ x.T
print("OK")
'''
    try:
        result = subprocess.run(
            ["python", "-c", test_script],
            capture_output=True,
            text=True,
            timeout=30,
            env={"CUDA_VISIBLE_DEVICES": str(gpu_id)},
        )
        return "OK" in result.stdout
    except Exception:
        return False


def load_state() -> dict[str, JobState]:
    """Load sweep state from file."""
    if STATE_FILE.exists():
        with open(STATE_FILE) as f:
            data = json.load(f)
            return {k: JobState(**v) for k, v in data.items()}
    return {}


def save_state(state: dict[str, JobState]) -> None:
    """Save sweep state to file."""
    with open(STATE_FILE, "w") as f:
        json.dump({k: asdict(v) for k, v in state.items()}, f, indent=2)


def launch_job(config_name: str, gpu_id: int) -> tuple[int, str]:
    """Launch a training job on specified GPU."""
    LOGS_DIR.mkdir(exist_ok=True)

    log_file = LOGS_DIR / f"{config_name}_gpu{gpu_id}.log"
    config_file = CONFIGS_DIR / f"{config_name}.yaml"

    cmd = f"""
    CUDA_VISIBLE_DEVICES={gpu_id} nohup python run_exp.py \\
        --config_file {config_file} \\
        --offline \\
        > {log_file} 2>&1 &
    echo $!
    """

    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    pid = int(result.stdout.strip())

    return pid, str(log_file)


def check_job_status(pid: int) -> str:
    """Check if a job is still running."""
    try:
        result = subprocess.run(
            ["ps", "-p", str(pid), "-o", "state="],
            capture_output=True,
            text=True,
        )
        if result.returncode == 0 and result.stdout.strip():
            return "running"
        return "completed"  # Process ended
    except Exception:
        return "unknown"


def extract_metrics_from_log(log_file: str) -> dict:
    """Extract best metrics from training log."""
    metrics = {
        "best_frac_oracle": None,
        "best_top1_acc": None,
        "best_val_acc": None,
        "final_epoch": None,
    }

    try:
        with open(log_file) as f:
            content = f.read()

        # Find best frac_oracle
        import re
        frac_matches = re.findall(r"frac_oracle=(\d+\.\d+)", content)
        if frac_matches:
            metrics["best_frac_oracle"] = max(float(x) for x in frac_matches)

        # Find best top1_acc
        acc_matches = re.findall(r"top1_acc['\"]?: ?(\d+\.\d+)", content)
        if acc_matches:
            metrics["best_top1_acc"] = max(float(x) for x in acc_matches)

        # Find best val_acc
        val_matches = re.findall(r"val_acc=(\d+\.\d+)", content)
        if val_matches:
            metrics["best_val_acc"] = max(float(x) for x in val_matches)

        # Check for early stopping or completion
        if "Early stopping" in content:
            epoch_match = re.search(r"Early stopping at epoch (\d+)", content)
            if epoch_match:
                metrics["final_epoch"] = int(epoch_match.group(1))

    except Exception as e:
        print(f"Error extracting metrics from {log_file}: {e}")

    return metrics


def print_status(state: dict[str, JobState]) -> None:
    """Print current sweep status."""
    print("\n" + "=" * 80)
    print("QWEN3 SWEEP STATUS")
    print("=" * 80)

    for config_name in SWEEP_CONFIGS:
        job = state.get(config_name)
        if job is None:
            print(f"  {config_name}: NOT STARTED")
        else:
            # Update status if running
            if job.status == "running" and job.pid:
                job.status = check_job_status(job.pid)
                if job.status == "completed":
                    job.end_time = datetime.now().isoformat()
                    metrics = extract_metrics_from_log(job.log_file)
                    job.best_frac_oracle = metrics.get("best_frac_oracle")
                    job.best_top1_acc = metrics.get("best_top1_acc")
                    job.best_val_acc = metrics.get("best_val_acc")

            status_str = f"  {config_name}: {job.status.upper()}"
            if job.gpu_id is not None:
                status_str += f" (GPU {job.gpu_id})"
            if job.best_frac_oracle:
                status_str += f" | frac_oracle={job.best_frac_oracle:.3f}"
            if job.best_top1_acc:
                status_str += f" | top1_acc={job.best_top1_acc:.3f}"
            if job.best_val_acc:
                status_str += f" | val_acc={job.best_val_acc:.3f}"
            print(status_str)

    save_state(state)
    print("=" * 80)


def launch_sweep(state: dict[str, JobState]) -> None:
    """Launch all pending jobs on healthy GPUs."""
    print("\nChecking GPU health...")
    healthy_gpus = check_gpu_health()

    if not healthy_gpus:
        print("ERROR: No healthy GPUs available!")
        return

    print(f"\nHealthy GPUs: {healthy_gpus}")

    # Find jobs that need to run
    pending_configs = []
    for config_name in SWEEP_CONFIGS:
        job = state.get(config_name)
        if job is None or job.status in ("pending", "failed"):
            pending_configs.append(config_name)

    if not pending_configs:
        print("All jobs already running or completed!")
        return

    print(f"Jobs to launch: {len(pending_configs)}")

    # Assign jobs to GPUs
    for i, config_name in enumerate(pending_configs):
        if i >= len(healthy_gpus):
            print(f"  {config_name}: QUEUED (no available GPU)")
            state[config_name] = JobState(
                config_name=config_name,
                gpu_id=-1,
                pid=None,
                status="pending",
                start_time=None,
                end_time=None,
                log_file="",
                best_frac_oracle=None,
                best_top1_acc=None,
                best_val_acc=None,
            )
            continue

        gpu_id = healthy_gpus[i]
        print(f"  Launching {config_name} on GPU {gpu_id}...")

        pid, log_file = launch_job(config_name, gpu_id)

        state[config_name] = JobState(
            config_name=config_name,
            gpu_id=gpu_id,
            pid=pid,
            status="running",
            start_time=datetime.now().isoformat(),
            end_time=None,
            log_file=log_file,
            best_frac_oracle=None,
            best_top1_acc=None,
            best_val_acc=None,
        )

        time.sleep(2)  # Brief pause between launches

    save_state(state)
    print("\nSweep launched! Use --status to monitor progress.")


def export_results(state: dict[str, JobState]) -> None:
    """Export final results to CSV."""
    import csv

    output_file = "qwen3_sweep_results.csv"

    with open(output_file, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "config", "loss_type", "learning_rate", "difficulty_head",
            "frac_oracle", "top1_acc", "val_acc", "status", "gpu", "duration"
        ])

        for config_name in SWEEP_CONFIGS:
            job = state.get(config_name)
            if job is None:
                continue

            # Parse config name for details
            parts = config_name.split("_")
            loss_type = parts[2] if len(parts) > 2 else "unknown"
            lr = parts[3] if len(parts) > 3 else "unknown"
            difficulty = "no" if "no_difficulty" in config_name else "yes"

            writer.writerow([
                config_name,
                loss_type,
                lr,
                difficulty,
                job.best_frac_oracle or "",
                job.best_top1_acc or "",
                job.best_val_acc or "",
                job.status,
                job.gpu_id,
                "",  # Duration would need calculation
            ])

    print(f"Results exported to {output_file}")


def main():
    parser = argparse.ArgumentParser(description="Qwen3 Sweep Launcher")
    parser.add_argument("--check-gpus", action="store_true", help="Check GPU health")
    parser.add_argument("--launch", action="store_true", help="Launch sweep")
    parser.add_argument("--status", action="store_true", help="Check status")
    parser.add_argument("--resume", action="store_true", help="Resume failed jobs")
    parser.add_argument("--export", action="store_true", help="Export results to CSV")

    args = parser.parse_args()
    state = load_state()

    if args.check_gpus:
        print("Checking GPU health...")
        healthy = check_gpu_health()
        print(f"\nHealthy GPUs: {healthy}")

    elif args.launch:
        launch_sweep(state)

    elif args.status:
        print_status(state)

    elif args.resume:
        # Mark failed jobs as pending and relaunch
        for config_name, job in state.items():
            if job.status == "failed":
                job.status = "pending"
        save_state(state)
        launch_sweep(state)

    elif args.export:
        export_results(state)

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
