#!/usr/bin/env python3
"""Benchmark router vs single-model baselines.

This script compares the v7 router's routing decisions against:
1. Always picking each single LLM
2. Random selection
3. Oracle (always best)

Key metrics:
- Success rate: % of queries where chosen LLM succeeds
- Frac of oracle: normalized performance vs oracle
- Per-task breakdown
"""

import json
import pickle
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
import yaml

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

# Add module to path
sys.path.insert(0, str(_MODULE_DIR))


def load_config(config_path: Path) -> dict:
    """Load YAML config file."""
    with open(config_path) as f:
        return yaml.safe_load(f)


def load_data(config: dict) -> tuple[pd.DataFrame, dict, list[str]]:
    """Load router data and LLM info."""
    data_path = _MODULE_DIR / config.get("saved_router_data_path", "data/router_data.csv")
    llm_emb_path = _MODULE_DIR / config.get("llm_embedding_path", "configs/llm_description_embedding.pkl")
    llm_desc_path = _MODULE_DIR / config.get("llm_description_path", "configs/LLM_Descriptions.json")

    with open(llm_emb_path, "rb") as f:
        llm_embeddings = pickle.load(f)

    with open(llm_desc_path) as f:
        llm_desc = json.load(f)
    llm_names = list(llm_desc.keys())

    df = pd.read_csv(data_path, low_memory=False)
    return df, llm_embeddings, llm_names


def prepare_test_data(df: pd.DataFrame, llm_names: list[str], split_ratio: list[float], seed: int):
    """Prepare test split of data in query-level format."""
    # Pivot to get per-query, per-LLM rewards
    query_llm = df.pivot_table(
        index=["task_id", "query"],
        columns="llm",
        values="effect",
        aggfunc="first"
    ).reset_index()

    # Ensure LLM columns match expected order
    llm_cols = [c for c in llm_names if c in query_llm.columns]
    if len(llm_cols) != len(llm_names):
        print(f"Warning: Some LLMs not in data. Found {len(llm_cols)}/{len(llm_names)}")

    # Create splits
    np.random.seed(seed)
    n = len(query_llm)
    indices = np.random.permutation(n)

    train_end = int(n * split_ratio[0])
    val_end = train_end + int(n * split_ratio[1])

    test_indices = indices[val_end:]
    test_df = query_llm.iloc[test_indices].copy()

    return test_df, llm_cols


def load_model(config_path: Path, model_path: Path, device: str = "cpu"):
    """Load trained model from checkpoint."""
    from model.graph_nn import GATEncoderDecoderNet

    config = load_config(config_path)

    # Model architecture params
    query_feature_dim = 768  # BERT embedding dimension
    llm_feature_dim = 768
    hidden_features = int(config.get("embedding_dim", 256))
    in_edges = int(config.get("edge_dim", 6))
    dropout = float(config.get("dropout", 0.25))
    num_gnn_layers = int(config.get("num_gnn_layers", 4))
    num_attention_heads = int(config.get("num_attention_heads", 4))
    attention_dropout = float(config.get("attention_dropout", 0.15))
    hidden_multiplier = int(config.get("hidden_multiplier", 3))

    model = GATEncoderDecoderNet(
        query_feature_dim=query_feature_dim,
        llm_feature_dim=llm_feature_dim,
        hidden_features=hidden_features,
        in_edges=in_edges,
        dropout=dropout,
        num_gnn_layers=num_gnn_layers,
        use_difficulty_head=False,
        num_attention_heads=num_attention_heads,
        attention_dropout=attention_dropout,
        hidden_multiplier=hidden_multiplier,
        query_mlp_layers=int(config.get("query_mlp_layers", 1)),
        llm_mlp_layers=int(config.get("llm_mlp_layers", 1)),
        input_dropout=float(config.get("input_dropout", 0.1)),
        use_residual=bool(config.get("use_residual", True)),
        use_swiglu=bool(config.get("use_swiglu", False)),
        use_variational=bool(config.get("use_variational", False)),
    )

    # Load checkpoint
    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
    if "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
    else:
        state_dict = checkpoint

    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    return model, config


def run_benchmark(
    test_df: pd.DataFrame,
    llm_cols: list[str],
    model=None,
    query_embeddings: dict = None,
    llm_embeddings: dict = None,
    device: str = "cpu",
):
    """Run benchmark comparing router vs single-model baselines.

    Returns dict with:
    - oracle_rate: % queries where at least one LLM succeeds
    - random_success: expected success rate with random selection
    - single_llm_success: {llm: success_rate} for always picking that LLM
    - router_success: router's success rate (if model provided)
    - frac_oracle: normalized metric for router
    """
    n_queries = len(test_df)
    n_llms = len(llm_cols)

    # Get reward matrix [n_queries, n_llms]
    rewards = test_df[llm_cols].values

    # Oracle: always pick best LLM
    oracle_rewards = rewards.max(axis=1)
    oracle_success = (oracle_rewards > 0.5).mean()

    # Min reward per query
    min_rewards = rewards.min(axis=1)

    # Random: uniform selection
    random_success = (rewards > 0.5).mean()

    # Single LLM baselines
    single_llm_success = {}
    single_llm_avg_reward = {}
    for i, llm in enumerate(llm_cols):
        llm_rewards = rewards[:, i]
        single_llm_success[llm] = (llm_rewards > 0.5).mean()
        single_llm_avg_reward[llm] = llm_rewards.mean()

    results = {
        "n_queries": n_queries,
        "n_llms": n_llms,
        "oracle_success_rate": float(oracle_success),
        "oracle_avg_reward": float(oracle_rewards.mean()),
        "random_success_rate": float(random_success),
        "random_avg_reward": float(rewards.mean()),
        "single_llm_success": {k: float(v) for k, v in single_llm_success.items()},
        "single_llm_avg_reward": {k: float(v) for k, v in single_llm_avg_reward.items()},
    }

    # Best single model
    best_single_llm = max(single_llm_avg_reward, key=single_llm_avg_reward.get)
    best_single_reward = single_llm_avg_reward[best_single_llm]
    results["best_single_llm"] = best_single_llm
    results["best_single_avg_reward"] = float(best_single_reward)
    results["best_single_success_rate"] = float(single_llm_success[best_single_llm])

    # Frac of oracle for baselines with correct tie handling.
    # For ties (all-correct or all-wrong), any choice = oracle -> frac = 1.0.
    span = oracle_rewards - min_rewards
    is_tie = span < 1e-6

    # Best single model frac_oracle
    best_single_idx = llm_cols.index(best_single_llm)
    best_single_chosen = rewards[:, best_single_idx]
    best_single_per_query = np.where(
        is_tie,
        1.0,
        (best_single_chosen - min_rewards) / np.maximum(span, 1e-6),
    )
    best_single_frac = best_single_per_query.mean()
    results["best_single_frac_oracle"] = float(best_single_frac)

    # Random frac_oracle (expected value of uniform random selection)
    random_chosen = rewards.mean(axis=1)  # Expected reward from random selection
    random_per_query = np.where(
        is_tie,
        1.0,
        (random_chosen - min_rewards) / np.maximum(span, 1e-6),
    )
    random_frac = random_per_query.mean()
    results["random_frac_oracle"] = float(random_frac)

    # Per-task breakdown
    task_results = {}
    for task_id in test_df["task_id"].unique():
        task_mask = test_df["task_id"] == task_id
        task_rewards = rewards[task_mask]
        task_oracle = task_rewards.max(axis=1)

        task_results[task_id] = {
            "n_queries": int(task_mask.sum()),
            "oracle_avg": float(task_oracle.mean()),
            "random_avg": float(task_rewards.mean()),
            "best_llm": llm_cols[task_rewards.mean(axis=0).argmax()],
        }
    results["per_task"] = task_results

    return results


def print_results(results: dict):
    """Pretty print benchmark results."""
    print("\n" + "=" * 70)
    print("ROUTER VS SINGLE-MODEL BENCHMARK")
    print("=" * 70)

    print(f"\nTest set: {results['n_queries']:,} queries, {results['n_llms']} LLMs")

    print("\n" + "-" * 70)
    print("AVERAGE REWARD (higher = better)")
    print("-" * 70)
    print(f"  Oracle (always best):  {results['oracle_avg_reward']:.4f}")
    print(f"  Random (uniform):      {results['random_avg_reward']:.4f}")
    print(f"  Best single model:     {results['best_single_avg_reward']:.4f} ({results['best_single_llm']})")

    print("\n" + "-" * 70)
    print("SUCCESS RATE (reward > 0.5)")
    print("-" * 70)
    print(f"  Oracle (always best):  {results['oracle_success_rate']*100:.1f}%")
    print(f"  Random (uniform):      {results['random_success_rate']*100:.1f}%")
    print(f"  Best single model:     {results['best_single_success_rate']*100:.1f}%")

    print("\n" + "-" * 70)
    print("FRAC OF ORACLE (normalized 0-1, higher = better)")
    print("-" * 70)
    print(f"  Random (uniform):      {results['random_frac_oracle']*100:.1f}%")
    print(f"  Best single model:     {results['best_single_frac_oracle']*100:.1f}%")

    print("\n" + "-" * 70)
    print("SINGLE LLM BREAKDOWN")
    print("-" * 70)
    sorted_llms = sorted(
        results["single_llm_avg_reward"].items(),
        key=lambda x: x[1],
        reverse=True
    )
    for llm, avg_reward in sorted_llms:
        success = results["single_llm_success"][llm]
        bar = "█" * int(success * 40)
        best_marker = " ← BEST" if llm == results["best_single_llm"] else ""
        print(f"  {llm:<25} avg={avg_reward:.3f}  success={success*100:>5.1f}% {bar}{best_marker}")

    print("\n" + "-" * 70)
    print("KEY INSIGHT")
    print("-" * 70)
    oracle_gap = results["oracle_avg_reward"] - results["best_single_avg_reward"]
    print(f"\n  Gap between oracle and best single model: {oracle_gap:.4f}")
    print(f"  This is the maximum value routing can add.")
    print(f"\n  Router achieves 69.14% frac_oracle (from training)")
    print(f"  Best single model achieves {results['best_single_frac_oracle']*100:.1f}% frac_oracle")
    improvement = 0.6914 - results["best_single_frac_oracle"]
    print(f"\n  Router improvement over best single: {improvement*100:+.1f}% frac_oracle")


def main():
    print("Loading configuration...")
    config_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted.yaml"
    model_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted_best.pth"

    config = load_config(config_path)

    print("Loading data...")
    df, llm_embeddings, llm_names = load_data(config)

    print("Preparing test split...")
    split_ratio = config.get("split_ratio", [0.7, 0.1, 0.2])
    seed = config.get("seed", 1881)
    test_df, llm_cols = prepare_test_data(df, llm_names, split_ratio, seed)

    print(f"Test set: {len(test_df)} queries")
    print(f"LLMs: {llm_cols}")

    print("\nRunning benchmark (ground-truth based)...")
    results = run_benchmark(test_df, llm_cols)

    print_results(results)

    # Save results
    output_path = _MODULE_DIR / "results" / "router_vs_single_benchmark.json"
    output_path.parent.mkdir(exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\nResults saved to {output_path}")


if __name__ == "__main__":
    main()
