#!/usr/bin/env python3
"""Hyperparameter sweep for EmberGraphRouter.

Runs parallel experiments on GPUs 1-7, keeping GPU 0 for ongoing training.
Focuses on anti-collapse mechanisms and architectural choices.
"""

import argparse
import logging
import os
import subprocess
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import yaml

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

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


# Priority configs to test - focus on VIB, anti-collapse, and architecture
# Each config is a delta from the base config
PRIORITY_CONFIGS = [
    # 1. Baseline + VIB (current config already has this)
    {
        "name": "baseline_vib",
        "use_variational": True,
        "vib_beta": 0.001,
    },
    # 2. VIB with stronger beta
    {
        "name": "vib_strong",
        "use_variational": True,
        "vib_beta": 0.01,
    },
    # 3. VIB with weaker beta
    {
        "name": "vib_weak",
        "use_variational": True,
        "vib_beta": 0.0001,
    },
    # 4. No VIB (ablation)
    {
        "name": "no_vib",
        "use_variational": False,
    },
    # 5. VIB + residual connections
    {
        "name": "vib_residual",
        "use_variational": True,
        "vib_beta": 0.001,
        "use_residual": True,
    },
    # 6. Stronger contrastive (anti-collapse)
    {
        "name": "strong_contrastive",
        "use_variational": True,
        "vib_beta": 0.001,
        "contrastive_weight": 0.2,
    },
    # 7. Stronger uniformity (anti-collapse)
    {
        "name": "strong_uniformity",
        "use_variational": True,
        "vib_beta": 0.001,
        "uniformity_weight": 0.1,
    },
    # 8. Stronger LLM balance
    {
        "name": "strong_llm_balance",
        "use_variational": True,
        "vib_beta": 0.001,
        "llm_balance_weight": 0.2,
    },
    # 9. All anti-collapse boosted
    {
        "name": "all_anticollapse_strong",
        "use_variational": True,
        "vib_beta": 0.001,
        "contrastive_weight": 0.15,
        "uniformity_weight": 0.075,
        "llm_balance_weight": 0.15,
        "quality_entropy_weight": 0.02,
    },
    # 10. Higher dropout
    {
        "name": "high_dropout",
        "use_variational": True,
        "vib_beta": 0.001,
        "dropout": 0.5,
    },
    # 11. Lower dropout
    {
        "name": "low_dropout",
        "use_variational": True,
        "vib_beta": 0.001,
        "dropout": 0.3,
    },
    # 12. 4-layer GNN (more capacity)
    {
        "name": "gnn_4layer",
        "use_variational": True,
        "vib_beta": 0.001,
        "num_gnn_layers": 4,
    },
    # 13. 2-layer GNN (less capacity)
    {
        "name": "gnn_2layer",
        "use_variational": True,
        "vib_beta": 0.001,
        "num_gnn_layers": 2,
    },
    # 14. Full combo: VIB + residual + strong anti-collapse
    {
        "name": "full_combo",
        "use_variational": True,
        "vib_beta": 0.001,
        "use_residual": True,
        "contrastive_weight": 0.15,
        "uniformity_weight": 0.075,
        "dropout": 0.45,
    },
]


def run_single_experiment(
    config_overrides: dict,
    gpu_id: int,
    base_config_path: str,
    train_epochs: int = 50,
    early_stop: int = 15,
) -> dict:
    """Run a single training experiment on specified GPU."""
    run_name = config_overrides.pop("name", f"run_gpu{gpu_id}")

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

    # Load base config
    with open(base_config_path, "r") as f:
        config = yaml.safe_load(f)

    # Apply overrides
    config.update(config_overrides)

    # Sweep-specific settings
    config["model_path"] = f"model_path/sweep_{run_name}.pth"
    config["train_epoch"] = train_epochs
    config["early_stop_patience"] = early_stop

    # Create temp config
    tmp_config_path = _MODULE_DIR / f"configs/sweep_{run_name}.yaml"
    with open(tmp_config_path, "w") as f:
        yaml.dump(config, f)

    cmd = [
        sys.executable,
        str(_MODULE_DIR / "run_exp.py"),
        "--config_file", str(tmp_config_path),
    ]

    logger.info(f"[GPU {gpu_id}] Starting {run_name}")

    start_time = time.time()
    result = subprocess.run(
        cmd,
        env=env,
        capture_output=True,
        text=True,
        cwd=str(_MODULE_DIR),
    )
    elapsed = time.time() - start_time

    # Parse metrics from output
    metrics = {
        "run_name": run_name,
        "gpu_id": gpu_id,
        "elapsed_min": elapsed / 60,
        "config": config_overrides,
    }

    import re

    # Find best score
    best_scores = re.findall(r"best model saved \(score=([\d.]+)\)", result.stdout, re.I)
    if best_scores:
        metrics["best_score"] = float(best_scores[-1])

    # Find final frac_oracle
    frac_oracles = re.findall(r"frac_of_oracle['\"]?:\s*([\d.]+)", result.stdout)
    if frac_oracles:
        metrics["frac_oracle"] = float(frac_oracles[-1])

    # Find final val_acc
    val_accs = re.findall(r"val_acc=([\d.]+)", result.stdout)
    if val_accs:
        metrics["val_acc"] = float(val_accs[-1])

    metrics["status"] = "success" if result.returncode == 0 else "failed"

    if result.returncode != 0:
        logger.error(f"[GPU {gpu_id}] {run_name} FAILED")
        # Save error log
        with open(_MODULE_DIR / f"logs/sweep_{run_name}_error.log", "w") as f:
            f.write(result.stderr)
    else:
        logger.info(
            f"[GPU {gpu_id}] {run_name} done in {elapsed/60:.1f}min | "
            f"score={metrics.get('best_score', 'N/A'):.4f} | "
            f"frac_oracle={metrics.get('frac_oracle', 'N/A')}"
        )

    return metrics


def main():
    parser = argparse.ArgumentParser(description="Run hyperparameter sweep")
    parser.add_argument("--gpus", type=str, default="1,2,3,4,5,6,7",
                        help="Comma-separated GPU IDs to use")
    parser.add_argument("--base_config", type=str,
                        default=str(_MODULE_DIR / "configs/best_listwise_gat3_20251213.yaml"),
                        help="Base config to modify")
    parser.add_argument("--epochs", type=int, default=50,
                        help="Training epochs per run")
    parser.add_argument("--early_stop", type=int, default=15,
                        help="Early stopping patience")
    parser.add_argument("--max_runs", type=int, default=14,
                        help="Maximum number of experiments")
    args = parser.parse_args()

    gpu_ids = [int(g) for g in args.gpus.split(",")]
    logger.info(f"Using GPUs: {gpu_ids} ({len(gpu_ids)} available)")

    # Ensure logs directory exists
    (_MODULE_DIR / "logs").mkdir(exist_ok=True)
    (_MODULE_DIR / "model_path").mkdir(exist_ok=True)

    configs = PRIORITY_CONFIGS[:args.max_runs]
    logger.info(f"Running {len(configs)} experiments ({args.epochs} epochs each)")

    results = []
    config_queue = list(configs)

    with ProcessPoolExecutor(max_workers=len(gpu_ids)) as executor:
        futures = {}
        gpu_available = list(gpu_ids)

        while config_queue or futures:
            # Submit jobs for available GPUs
            while config_queue and gpu_available:
                config = config_queue.pop(0).copy()  # Copy to avoid mutation
                gpu = gpu_available.pop(0)
                future = executor.submit(
                    run_single_experiment,
                    config, gpu, args.base_config,
                    args.epochs, args.early_stop,
                )
                futures[future] = gpu

            # Wait for completions
            if futures:
                done = next(as_completed(futures))
                gpu = futures.pop(done)
                gpu_available.append(gpu)

                try:
                    result = done.result()
                    results.append(result)
                except Exception as e:
                    logger.error(f"GPU {gpu} job failed: {e}")

    # Summary
    print("\n" + "="*70)
    print("SWEEP RESULTS")
    print("="*70)

    successful = [r for r in results if r.get("status") == "success"]
    successful.sort(key=lambda x: x.get("best_score", 0), reverse=True)

    print(f"\n{'Rank':<5} {'Name':<25} {'Score':<8} {'FracOracle':<12} {'Time':<8}")
    print("-"*70)

    for i, r in enumerate(successful, 1):
        print(
            f"{i:<5} {r['run_name']:<25} "
            f"{r.get('best_score', 0):<8.4f} "
            f"{r.get('frac_oracle', 0):<12.4f} "
            f"{r.get('elapsed_min', 0):<8.1f}m"
        )

    if successful:
        best = successful[0]
        print(f"\n{'='*70}")
        print(f"BEST CONFIG: {best['run_name']}")
        print(f"Score: {best.get('best_score', 'N/A')}")
        print(f"Config overrides: {best['config']}")

    # Save results
    results_path = _MODULE_DIR / "sweep_results.yaml"
    with open(results_path, "w") as f:
        yaml.dump(results, f, default_flow_style=False)
    print(f"\nResults saved to {results_path}")


if __name__ == "__main__":
    main()
