#!/usr/bin/env python3
"""Launch instruction-conditioned embedding comparison sweep.

This script runs a controlled experiment comparing:
- Baseline: Qwen3 with default retrieval instruction
- Treatment: Qwen3 with router-specific instruction

The comparison is run at multiple embedding dimensions (256, 512, 1024)
to understand how instruction conditioning interacts with dimensionality.

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

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

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

# Experiment configurations
EXPERIMENTS = [
    # (config_name, description)
    ("instruct_baseline_256", "Baseline: retrieval instruction, 256-dim"),
    ("instruct_router_v1_256", "Treatment: router instruction, 256-dim"),
    ("instruct_baseline_512", "Baseline: retrieval instruction, 512-dim"),
    ("instruct_router_v1_512", "Treatment: router instruction, 512-dim"),
]


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

    # Required files
    required_files = [
        # Baseline data
        ("data/router_data_qwen3.parquet", "Baseline parquet"),
        ("configs/llm_embedding_qwen3_256.pkl", "Baseline LLM embeddings 256"),
        ("configs/llm_embedding_qwen3_512.pkl", "Baseline LLM embeddings 512"),
        # Treatment data
        ("data/router_data_qwen3_instruct_router_v1.parquet", "Treatment parquet"),
        ("configs/llm_embedding_qwen3_instruct_router_v1_256.pkl", "Treatment LLM embeddings 256"),
        ("configs/llm_embedding_qwen3_instruct_router_v1_512.pkl", "Treatment LLM embeddings 512"),
    ]

    all_present = True
    for path, desc in required_files:
        full_path = _MODULE_DIR / path
        exists = full_path.exists()
        status = "✓" if exists else "✗"
        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

    if not all_present:
        print("\n⚠️  Some files are missing!")
        print("   Run generate_qwen3_instruct_embeddings.py first to create treatment data.")
    else:
        print("\n✓ 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)

    cmd = [
        "python", "run_exp.py",
        "--config", 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 instruction comparison sweep")
    print("=" * 70)

    # Check GPU availability
    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}")
        gpus = []

    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, log_path))
        time.sleep(1)  # Stagger launches slightly

    # Save PIDs for status checking
    pids_file = _MODULE_DIR / "logs" / "instruct_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_instruct_sweep.py --status")
    print("  tail -f logs/instruct_*.log")
    print("=" * 70)


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

    pids_file = _MODULE_DIR / "logs" / "instruct_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", "3", 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("Experiment 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

        # Parse log for final metrics
        with open(log_path) as f:
            log_content = f.read()

        # Look for best validation metrics
        # Example: "Best val loss: 0.1234, Best val acc: 0.5678"
        metrics = {}
        for line in log_content.split("\n"):
            if "Best val" in line or "Test" in line or "final" in line.lower():
                print(f"  {config_name}: {line.strip()}")
                break

    # Summary comparison
    print("\n" + "=" * 70)
    print("Summary: Baseline vs Treatment")
    print("=" * 70)

    # Group by dimension
    for dim in ["256", "512"]:
        print(f"\nDimension: {dim}")
        baseline = f"instruct_baseline_{dim}"
        treatment = f"instruct_router_v1_{dim}"
        print(f"  Baseline ({baseline}): [check logs]")
        print(f"  Treatment ({treatment}): [check logs]")


def main():
    parser = argparse.ArgumentParser(
        description="Instruction embedding 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()
