#!/usr/bin/env python3
"""Launch quality_only loss comparison sweep.

This script runs a controlled experiment comparing:
- BCE: Standard binary cross-entropy
- Focal: Down-weights easy examples (alpha=0.25, gamma=2.0)
- Brier: Proper scoring rule for calibration

All experiments use:
- Instruction-conditioned Qwen3 embeddings (router_data_qwen3_instruct_router_v1.parquet)
- 3-layer GAT with warmup + EMA
- Identical hyperparameters except loss function

Usage:
    cd ~/embergraphrouter
    python scripts/launch_quality_only_loss_sweep.py --check   # Check data availability
    python scripts/launch_quality_only_loss_sweep.py --launch  # Launch sweep
    python scripts/launch_quality_only_loss_sweep.py --status  # Check status
    python scripts/launch_quality_only_loss_sweep.py --results # Collect results
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time
from pathlib import Path

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

# Experiment configurations
EXPERIMENTS = [
    ("quality_only_loss_sweep_bce", "BCE loss (baseline)"),
    ("quality_only_loss_sweep_focal", "Focal loss (hard example focus)"),
    ("quality_only_loss_sweep_brier", "Brier loss (calibration)"),
]


def check_data_availability():
    """Check if required data files exist."""
    print("=" * 70)
    print("Checking data availability for quality_only loss sweep")
    print("=" * 70)

    required_files = [
        ("data/router_data_qwen3_instruct_router_v1.parquet", "Instruction-conditioned embeddings"),
        ("configs/llm_embedding_qwen3_instruct_router_v1_256.pkl", "LLM embeddings 256-dim"),
    ]

    all_present = True
    for path, desc in required_files:
        full_path = _MODULE_DIR / path
        exists = full_path.exists()
        status = "OK" if exists else "MISSING"
        size = ""
        if exists:
            size_bytes = full_path.stat().st_size
            if size_bytes > 1e9:
                size = f" ({size_bytes / 1e9:.1f} GB)"
            elif size_bytes > 1e6:
                size = f" ({size_bytes / 1e6:.1f} MB)"
            else:
                size = f" ({size_bytes / 1e3:.1f} KB)"
        print(f"  [{status}] {desc}: {path}{size}")
        if not exists:
            all_present = False

    # Check config files
    print("\nConfig files:")
    for config_name, desc in EXPERIMENTS:
        config_path = _MODULE_DIR / "configs" / f"{config_name}.yaml"
        exists = config_path.exists()
        status = "OK" if exists else "MISSING"
        print(f"  [{status}] {desc}: configs/{config_name}.yaml")
        if not exists:
            all_present = False

    if not all_present:
        print("\n[WARNING] Some files are missing!")
        print("   Generate instruction embeddings first if data is missing.")
    else:
        print("\n[OK] All required files present. Ready to launch sweep.")

    return all_present


def launch_experiment(config_name: str, gpu_id: int = 0):
    """Launch a single experiment on specified GPU."""
    config_path = _MODULE_DIR / "configs" / f"{config_name}.yaml"
    log_path = _MODULE_DIR / "logs" / f"{config_name}.log"

    log_path.parent.mkdir(exist_ok=True)

    # Use venv Python if available
    venv_python = _MODULE_DIR / ".venv" / "bin" / "python3"
    python_cmd = str(venv_python) if venv_python.exists() else "python3"

    cmd = [
        python_cmd, "run_exp.py",
        "--config_file", str(config_path),
    ]

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

    with open(log_path, "w") as log_file:
        process = subprocess.Popen(
            cmd,
            cwd=str(_MODULE_DIR),
            stdout=log_file,
            stderr=subprocess.STDOUT,
            env=env,
        )

    return process.pid, log_path


def launch_sweep():
    """Launch all experiments in parallel on different GPUs."""
    print("=" * 70)
    print("Launching quality_only loss comparison sweep")
    print("=" * 70)

    # Check GPU availability
    gpus = []
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,name,memory.free", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10
        )
        gpus = result.stdout.strip().split("\n")
        print(f"\nAvailable GPUs: {len(gpus)}")
        for gpu in gpus:
            print(f"  {gpu}")
    except Exception as e:
        print(f"Warning: Could not query GPUs: {e}")

    print("\nLaunching experiments:")
    pids = []
    for i, (config_name, desc) in enumerate(EXPERIMENTS):
        gpu_id = i % max(len(gpus), 1)
        pid, log_path = launch_experiment(config_name, gpu_id)
        print(f"  [{config_name}] GPU {gpu_id} -> PID {pid}")
        print(f"     Log: {log_path}")
        pids.append((config_name, pid, str(log_path)))
        time.sleep(1)  # Stagger launches

    # Save PIDs for status checking
    pids_file = _MODULE_DIR / "logs" / "quality_only_loss_sweep_pids.json"
    with open(pids_file, "w") as f:
        json.dump(pids, f, indent=2)
    print(f"\nPIDs saved to: {pids_file}")

    print("\n" + "=" * 70)
    print("Sweep launched! Monitor with:")
    print("  python scripts/launch_quality_only_loss_sweep.py --status")
    print("  tail -f logs/quality_only_loss_sweep_*.log")
    print("=" * 70)


def check_status():
    """Check status of running experiments."""
    print("=" * 70)
    print("Experiment Status")
    print("=" * 70)

    pids_file = _MODULE_DIR / "logs" / "quality_only_loss_sweep_pids.json"
    if not pids_file.exists():
        print("No sweep running. Launch with --launch first.")
        return

    with open(pids_file) as f:
        pids = json.load(f)

    for config_name, pid, log_path in pids:
        # Check if process is running
        try:
            os.kill(pid, 0)
            running = True
        except OSError:
            running = False

        status = "RUNNING" if running else "COMPLETED"

        # Get last few lines of log
        log_tail = ""
        if Path(log_path).exists():
            result = subprocess.run(
                ["tail", "-n", "5", log_path],
                capture_output=True, text=True
            )
            log_tail = result.stdout.strip()

        print(f"\n{config_name}:")
        print(f"  Status: {status}")
        print(f"  PID: {pid}")
        if log_tail:
            print(f"  Last log:")
            for line in log_tail.split("\n"):
                print(f"    {line}")


def collect_results():
    """Collect and compare results from completed experiments."""
    print("=" * 70)
    print("Quality-Only Loss Sweep Results")
    print("=" * 70)

    results = {}
    for config_name, desc in EXPERIMENTS:
        log_path = _MODULE_DIR / "logs" / f"{config_name}.log"
        if not log_path.exists():
            print(f"  {config_name}: No log file found")
            continue

        with open(log_path) as f:
            log_content = f.read()

        # Parse final metrics - look for test frac_of_oracle
        frac_oracle = None
        ece = None
        for line in log_content.split("\n"):
            # Look for test metrics
            if "frac_of_oracle" in line.lower() and "test" in line.lower():
                match = re.search(r"frac_of_oracle[=:\s]*([\d.]+)", line, re.IGNORECASE)
                if match:
                    frac_oracle = float(match.group(1))
            if "ece" in line.lower():
                match = re.search(r"ece[=:\s]*([\d.]+)", line, re.IGNORECASE)
                if match:
                    ece = float(match.group(1))

        results[config_name] = {
            "desc": desc,
            "frac_oracle": frac_oracle,
            "ece": ece,
        }
        print(f"\n{desc} ({config_name}):")
        print(f"  frac_of_oracle: {frac_oracle if frac_oracle else 'N/A'}")
        print(f"  ECE (calibration): {ece if ece else 'N/A'}")

    # Summary comparison
    print("\n" + "=" * 70)
    print("Summary Comparison")
    print("=" * 70)
    print(f"{'Loss':<15} {'frac_of_oracle':>15} {'ECE':>10}")
    print("-" * 42)
    for config_name, data in results.items():
        loss_name = config_name.split("_")[-1].upper()
        frac = f"{data['frac_oracle']:.4f}" if data['frac_oracle'] else "N/A"
        ece = f"{data['ece']:.4f}" if data['ece'] else "N/A"
        print(f"{loss_name:<15} {frac:>15} {ece:>10}")

    # Determine winner
    valid_results = [(k, v) for k, v in results.items() if v['frac_oracle'] is not None]
    if valid_results:
        best = max(valid_results, key=lambda x: x[1]['frac_oracle'])
        print(f"\nBest loss: {best[0].split('_')[-1].upper()} (frac_of_oracle={best[1]['frac_oracle']:.4f})")


def main():
    parser = argparse.ArgumentParser(
        description="Quality-only loss comparison sweep"
    )
    parser.add_argument("--check", action="store_true", help="Check data availability")
    parser.add_argument("--launch", action="store_true", help="Launch sweep")
    parser.add_argument("--status", action="store_true", help="Check experiment status")
    parser.add_argument("--results", action="store_true", help="Collect results")
    args = parser.parse_args()

    if args.check:
        check_data_availability()
    elif args.launch:
        if check_data_availability():
            launch_sweep()
    elif args.status:
        check_status()
    elif args.results:
        collect_results()
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
