#!/usr/bin/env python3
"""Evaluate intelligent routing strategies using P(correct | query, LLM) outputs.

The GNN outputs calibrated probabilities for each LLM, not just a routing decision.
This script evaluates multiple strategies for using these probabilities intelligently:

1. Argmax (baseline): Pick LLM with highest P(correct)
2. Entropy cascade: High entropy → premium model, low entropy → predicted
3. Expected utility: Maximize w_p * P(correct) - w_c * cost - w_l * latency
4. Confidence-gated: Use prediction if confident, else fallback to premium

Usage:
    python -m embergraphrouter.scripts.evaluate_intelligent_routing \
        --data data/router_data_qwen3.parquet
"""

import argparse
import random
import sys
from dataclasses import dataclass
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

sys.path.insert(0, str(_MODULE_DIR))
sys.path.insert(0, str(_MODULE_DIR.parent))


@dataclass
class RoutingResult:
    """Result of a routing strategy evaluation."""

    name: str
    accuracy: float  # % of queries where chosen LLM == oracle LLM
    frac_oracle: float  # (chosen_reward - min) / (oracle_reward - min)
    avg_reward: float
    avg_cost: float
    premium_fraction: float  # % of queries routed to premium model
    description: str


def compute_entropy(probs: np.ndarray) -> np.ndarray:
    """Compute normalized entropy of probability distribution.

    Args:
        probs: (n_queries, n_llms) probability matrix

    Returns:
        (n_queries,) normalized entropy values in [0, 1]
    """
    n_llms = probs.shape[1]
    entropy = -np.sum(probs * np.log(probs + 1e-8), axis=1)
    max_entropy = np.log(n_llms)
    return entropy / max_entropy


def evaluate_strategy(
    strategy_name: str,
    strategy_choices: np.ndarray,  # (n_queries,) LLM indices
    ground_truth: np.ndarray,  # (n_queries, n_llms) rewards
    costs: np.ndarray,  # (n_llms,) cost per LLM
    premium_idx: int,
    description: str = "",
) -> RoutingResult:
    """Evaluate a routing strategy.

    Args:
        strategy_name: Name for this strategy
        strategy_choices: (n_queries,) array of chosen LLM indices
        ground_truth: (n_queries, n_llms) reward matrix
        costs: (n_llms,) cost per LLM
        premium_idx: Index of the premium model
        description: Strategy description

    Returns:
        RoutingResult with metrics
    """
    n_queries = len(strategy_choices)

    # Compute oracle (best LLM per query)
    oracle_choices = ground_truth.argmax(axis=1)
    oracle_rewards = ground_truth.max(axis=1)
    min_rewards = ground_truth.min(axis=1)

    # Get rewards for our choices
    chosen_rewards = np.array([
        ground_truth[i, strategy_choices[i]] for i in range(n_queries)
    ])

    # Accuracy: exact match with oracle
    accuracy = (strategy_choices == oracle_choices).mean()

    # Frac oracle: normalized reward 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
    per_query_frac = np.where(
        is_tie,
        1.0,  # Ties: any choice matches oracle
        (chosen_rewards - min_rewards) / np.maximum(span, 1e-6),
    )
    frac_oracle = per_query_frac.mean()

    # Average reward
    avg_reward = chosen_rewards.mean()

    # Average cost
    avg_cost = np.mean([costs[strategy_choices[i]] for i in range(n_queries)])

    # Premium fraction
    premium_fraction = (strategy_choices == premium_idx).mean()

    return RoutingResult(
        name=strategy_name,
        accuracy=accuracy,
        frac_oracle=frac_oracle,
        avg_reward=avg_reward,
        avg_cost=avg_cost,
        premium_fraction=premium_fraction,
        description=description,
    )


def run_evaluation(
    predictions: np.ndarray,  # (n_queries, n_llms) logits or probs
    ground_truth: np.ndarray,  # (n_queries, n_llms) rewards
    llm_names: list[str],
    costs: np.ndarray,
    latencies: np.ndarray,
    perf_weight: float = 0.7,
    cost_weight: float = 0.2,
    latency_weight: float = 0.1,
) -> list[RoutingResult]:
    """Run all routing strategies and return results.

    Args:
        predictions: Model logits (n_queries, n_llms)
        ground_truth: Reward matrix (n_queries, n_llms)
        llm_names: List of LLM names
        costs: Per-LLM costs
        latencies: Per-LLM latencies
        perf_weight: Weight for performance in expected utility
        cost_weight: Weight for cost
        latency_weight: Weight for latency

    Returns:
        List of RoutingResult for each strategy
    """
    n_queries, n_llms = predictions.shape

    # Convert logits to probabilities
    probs = F.softmax(torch.tensor(predictions), dim=1).numpy()

    # Identify premium and cheap models by name
    premium_names = ["Claude Opus 4.5", "GPT-4.1"]
    cheap_names = ["GPT-5 Nano", "GPT-4.1 Mini"]

    premium_idx = None
    cheap_idx = None
    for i, name in enumerate(llm_names):
        if any(p in name for p in premium_names):
            premium_idx = i
            break
    for i, name in enumerate(llm_names):
        if any(c in name for c in cheap_names):
            cheap_idx = i
            break

    if premium_idx is None:
        premium_idx = costs.argmax()  # Most expensive as premium
    if cheap_idx is None:
        cheap_idx = costs.argmin()  # Cheapest as cheap

    print(f"Premium model: {llm_names[premium_idx]} (idx={premium_idx})")
    print(f"Cheap model: {llm_names[cheap_idx]} (idx={cheap_idx})")

    results = []

    # Compute prediction entropy
    entropy = compute_entropy(probs)
    max_prob = probs.max(axis=1)
    argmax_choices = probs.argmax(axis=1)

    # === BASELINES ===

    # 1. Oracle (always best)
    oracle_choices = ground_truth.argmax(axis=1)
    results.append(evaluate_strategy(
        "Oracle",
        oracle_choices,
        ground_truth,
        costs,
        premium_idx,
        "Always pick the best LLM (upper bound)",
    ))

    # 2. Random
    np.random.seed(1881)
    random_choices = np.random.randint(0, n_llms, size=n_queries)
    results.append(evaluate_strategy(
        "Random",
        random_choices,
        ground_truth,
        costs,
        premium_idx,
        "Uniform random selection",
    ))

    # 3. Best single model (by average reward)
    single_avg_rewards = ground_truth.mean(axis=0)
    best_single_idx = single_avg_rewards.argmax()
    always_best_single = np.full(n_queries, best_single_idx)
    results.append(evaluate_strategy(
        f"Always {llm_names[best_single_idx][:20]}",
        always_best_single,
        ground_truth,
        costs,
        premium_idx,
        f"Always pick {llm_names[best_single_idx]}",
    ))

    # 4. Always premium
    always_premium = np.full(n_queries, premium_idx)
    results.append(evaluate_strategy(
        f"Always Premium",
        always_premium,
        ground_truth,
        costs,
        premium_idx,
        f"Always pick {llm_names[premium_idx]}",
    ))

    # === MODEL-BASED STRATEGIES ===

    # 5. Simple argmax (baseline model usage)
    results.append(evaluate_strategy(
        "Argmax P(correct)",
        argmax_choices,
        ground_truth,
        costs,
        premium_idx,
        "Pick LLM with highest P(correct)",
    ))

    # 6. Expected utility maximization
    # utility[i,j] = perf_weight * P(correct|query_i, llm_j)
    #               - cost_weight * cost_j - latency_weight * latency_j
    utility = (
        perf_weight * probs
        - cost_weight * costs[np.newaxis, :]
        - latency_weight * latencies[np.newaxis, :]
    )
    utility_choices = utility.argmax(axis=1)
    results.append(evaluate_strategy(
        "Expected Utility",
        utility_choices,
        ground_truth,
        costs,
        premium_idx,
        f"Max(0.7*P(c) - 0.2*cost - 0.1*latency)",
    ))

    # 7. Entropy cascade at different thresholds
    for thresh in [0.3, 0.5, 0.7]:
        cascade_choices = np.where(
            entropy > thresh,
            premium_idx,  # High entropy -> premium
            argmax_choices,  # Low entropy -> model prediction
        )
        results.append(evaluate_strategy(
            f"Entropy Cascade (t={thresh})",
            cascade_choices,
            ground_truth,
            costs,
            premium_idx,
            f"High entropy (>{thresh}) -> premium, else -> argmax",
        ))

    # 8. Confidence-gated (use prediction only if confident)
    for conf_thresh in [0.5, 0.7, 0.9]:
        gated_choices = np.where(
            max_prob > conf_thresh,
            argmax_choices,  # Confident -> use prediction
            premium_idx,  # Not confident -> premium
        )
        results.append(evaluate_strategy(
            f"Confidence Gate (>{conf_thresh})",
            gated_choices,
            ground_truth,
            costs,
            premium_idx,
            f"If max P(c)>{conf_thresh} -> argmax, else -> premium",
        ))

    # 9. Cost-aware cascade: cheap for confident, premium for uncertain
    for thresh in [0.5, 0.7]:
        cost_cascade_choices = np.where(
            max_prob > thresh,
            cheap_idx,  # Confident -> cheap model
            premium_idx,  # Uncertain -> premium
        )
        results.append(evaluate_strategy(
            f"Cost Cascade (t={thresh})",
            cost_cascade_choices,
            ground_truth,
            costs,
            premium_idx,
            f"If confident (>{thresh}) -> cheap, else -> premium",
        ))

    # 10. Weighted by P(correct) ranking
    # For each LLM, compute E[reward] = P(correct) * effect
    # This uses the actual prediction to weight expected value

    return results


def print_results(results: list[RoutingResult]):
    """Print formatted results table."""
    print("\n" + "=" * 100)
    print("INTELLIGENT ROUTING STRATEGY EVALUATION")
    print("=" * 100)

    print(f"\n{'Strategy':<30} {'Accuracy':>10} {'FracOracle':>12} {'AvgReward':>12} {'Premium%':>10}")
    print("-" * 100)

    # Sort by frac_oracle
    sorted_results = sorted(results, key=lambda r: -r.frac_oracle)

    for r in sorted_results:
        is_baseline = r.name in ["Oracle", "Random"] or "Always" in r.name
        marker = "" if is_baseline else " ★" if r.name == "Argmax P(correct)" else ""
        print(
            f"{r.name:<30} "
            f"{r.accuracy*100:>9.1f}% "
            f"{r.frac_oracle*100:>11.1f}% "
            f"{r.avg_reward:>11.4f} "
            f"{r.premium_fraction*100:>9.1f}%"
            f"{marker}"
        )

    print("\n" + "-" * 100)
    print("KEY INSIGHTS:")
    print("-" * 100)

    # Find best model-based strategy
    model_strategies = [r for r in results if r.name not in ["Oracle", "Random"] and "Always" not in r.name]
    best_model = max(model_strategies, key=lambda r: r.frac_oracle)

    # Find best single model baseline
    baselines = [r for r in results if "Always" in r.name]
    best_baseline = max(baselines, key=lambda r: r.frac_oracle)

    oracle = next(r for r in results if r.name == "Oracle")
    argmax = next(r for r in results if r.name == "Argmax P(correct)")

    print(f"  Oracle:                  {oracle.frac_oracle*100:.1f}% frac_oracle")
    print(f"  Best single model:       {best_baseline.frac_oracle*100:.1f}% ({best_baseline.name})")
    print(f"  Simple argmax:           {argmax.frac_oracle*100:.1f}%")
    print(f"  Best intelligent routing: {best_model.frac_oracle*100:.1f}% ({best_model.name})")
    print()

    improvement = best_model.frac_oracle - best_baseline.frac_oracle
    if improvement > 0:
        print(f"  ✓ Best strategy improves over single model by {improvement*100:+.1f}%")
    else:
        print(f"  ✗ No improvement over single model baseline")

    improvement_over_argmax = best_model.frac_oracle - argmax.frac_oracle
    if best_model.name != "Argmax P(correct)" and improvement_over_argmax > 0:
        print(f"  ✓ Using entropy/confidence improves over simple argmax by {improvement_over_argmax*100:+.1f}%")


def main():
    parser = argparse.ArgumentParser(description="Evaluate intelligent routing strategies")
    parser.add_argument(
        "--data",
        type=Path,
        default=_MODULE_DIR / "data/router_data_qwen3.parquet",
        help="Router data path",
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=_MODULE_DIR / "configs/config.yaml",
        help="Config file path",
    )
    parser.add_argument("--seed", type=int, default=1881)
    args = parser.parse_args()

    print("Loading data...")

    # Load config
    if args.config.exists():
        with open(args.config) as f:
            config = yaml.safe_load(f)
    else:
        config = {"split_ratio": [0.7, 0.1, 0.2], "seed": 1881}

    split_ratio = config.get("split_ratio", [0.7, 0.1, 0.2])
    seed = config.get("seed", args.seed)

    # Load data - support both parquet and CSV
    data_path = str(args.data)
    if data_path.endswith(".csv"):
        df = pd.read_csv(data_path, low_memory=False)
    else:
        try:
            df = pd.read_parquet(
                data_path,
                columns=["query", "task_id", "llm", "effect", "cost", "latency"],
            )
        except Exception:
            df = pd.read_parquet(
                data_path,
                columns=["query", "task_id", "llm", "effect"],
            )

    # Add default cost/latency if not available
    if "cost" not in df.columns:
        df["cost"] = 0.0
    if "latency" not in df.columns:
        df["latency"] = 0.0

    print(f"Loaded {len(df)} rows, {df['query'].nunique()} unique queries")

    # Get LLM info
    llm_names = df["llm"].unique().tolist()[:7]  # First 7 LLMs
    n_llms = len(llm_names)
    print(f"LLMs: {llm_names}")

    # Filter to 7 LLMs
    df = df[df["llm"].isin(llm_names)]

    # Build (query, task_id) groups
    df["group_key"] = df["query"].astype(str) + "|||" + df["task_id"].astype(str)
    group_sizes = df.groupby("group_key").size()
    valid_groups = group_sizes[group_sizes == n_llms].index
    df = df[df["group_key"].isin(valid_groups)]

    # Create ground truth and cost matrices
    llm_to_idx = {name: i for i, name in enumerate(llm_names)}

    ground_truth_list = []
    cost_per_query_list = []
    latency_per_query_list = []
    task_ids_list = []

    for group_key, group_df in df.groupby("group_key", sort=False):
        gt_row = np.zeros(n_llms)
        cost_row = np.zeros(n_llms)
        latency_row = np.zeros(n_llms)

        for _, row in group_df.iterrows():
            llm_idx = llm_to_idx.get(row["llm"])
            if llm_idx is not None:
                gt_row[llm_idx] = row["effect"]
                cost_row[llm_idx] = row["cost"]
                latency_row[llm_idx] = row["latency"]

        ground_truth_list.append(gt_row)
        cost_per_query_list.append(cost_row)
        latency_per_query_list.append(latency_row)
        task_ids_list.append(group_df["task_id"].iloc[0])

    ground_truth = np.array(ground_truth_list)
    costs_per_query = np.array(cost_per_query_list)
    latencies_per_query = np.array(latency_per_query_list)
    task_ids = np.array(task_ids_list)

    # Get average costs/latencies per LLM
    costs = costs_per_query.mean(axis=0)
    latencies = latencies_per_query.mean(axis=0)

    print(f"Query groups: {len(ground_truth)}")

    # === STRATIFIED VALIDATION SPLIT ===
    grouped = df.drop_duplicates(subset=["group_key"])[["group_key", "task_id"]]
    task_to_indices = {}
    for i, row in grouped.reset_index(drop=True).iterrows():
        task_id = row["task_id"]
        task_to_indices.setdefault(task_id, []).append(i)

    rng = random.Random(seed)
    val_indices = []

    for task_id, indices in task_to_indices.items():
        indices_copy = indices.copy()
        rng.shuffle(indices_copy)
        n = len(indices_copy)
        train_c = int(n * split_ratio[0])
        val_c = int(n * split_ratio[1])
        val_indices.extend(indices_copy[train_c:train_c + val_c])

    print(f"Validation set: {len(val_indices)} query groups (stratified)")

    # Extract validation data
    val_ground_truth = ground_truth[val_indices]
    val_costs = costs_per_query[val_indices]
    val_latencies = latencies_per_query[val_indices]

    # === TRY LOADING ACTUAL MODEL ===
    actual_logits = None
    model_source = "simulated"

    try:
        from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

        # Look for config file
        config_candidates = [
            _MODULE_DIR / "configs/config.yaml",
            _MODULE_DIR / "configs/hybrid_lambda_rank_openai512.yaml",
        ]

        trainer_config_path = None
        for cp in config_candidates:
            if cp.exists():
                trainer_config_path = cp
                break

        if trainer_config_path:
            with open(trainer_config_path) as f:
                trainer_config = yaml.safe_load(f)

            # Initialize trainer
            trainer = GraphRouterTrainer(
                router_data_path=str(args.data),
                llm_path=str(_MODULE_DIR / "configs/LLM_Descriptions.json"),
                llm_embedding_path=str(_MODULE_DIR / "configs/llm_description_embedding.pkl"),
                config=trainer_config,
                wandb=None,
            )

            # Try loading best checkpoint
            checkpoint_candidates = [
                _MODULE_DIR / "model_path/hybrid_lambda_rank_openai512_best.pth",
                _MODULE_DIR / "model_path/best_model.pth",
                _MODULE_DIR / "registry/eb1-gnn-preview-12-19/model.pth",
            ]

            checkpoint_path = None
            for cp in checkpoint_candidates:
                if cp.exists():
                    checkpoint_path = cp
                    break

            if checkpoint_path:
                print(f"\n✓ Loading actual model from: {checkpoint_path}")
                state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
                trainer.gnn.model.load_state_dict(state_dict)
                trainer.gnn.model.eval()

                # Get predictions on validation set
                val_gnn_indices = trainer.gnn.val_indices
                print(f"  Model validation set: {len(val_gnn_indices)} queries")

                actual_logits_list = []
                with torch.no_grad():
                    for idx in val_gnn_indices:
                        pred_logits = trainer.gnn.predict_single(idx)
                        actual_logits_list.append(pred_logits)

                actual_logits = np.array(actual_logits_list)
                model_source = f"actual ({checkpoint_path.name})"
                print(f"  Got predictions: shape {actual_logits.shape}")

                # Update val_ground_truth to match GNN's validation indices
                val_ground_truth = trainer.gnn.labels[val_gnn_indices].cpu().numpy()
                print(f"  Updated ground truth: shape {val_ground_truth.shape}")
            else:
                print("\n⚠️  No checkpoint found, using simulated predictions")
        else:
            print("\n⚠️  No config found, using simulated predictions")

    except Exception as e:
        print(f"\n⚠️  Could not load actual model: {e}")
        print("   Using simulated predictions instead")

    # === USE ACTUAL OR SIMULATED PREDICTIONS ===
    if actual_logits is not None:
        print("\n" + "=" * 100)
        print(f"USING ACTUAL MODEL PREDICTIONS ({model_source})")
        print("=" * 100)
        predictions = actual_logits
    else:
        # Simulate predictions with varying confidence levels
        print("\n" + "=" * 100)
        print("SIMULATING MODEL PREDICTIONS (based on ground truth with noise)")
        print("=" * 100)
        print("\n⚠️  Note: Using simulated predictions for strategy evaluation.")
        print("   In production, load actual model with GraphAwareRouter.from_registry()")

        np.random.seed(seed)

        # Create simulated logits that produce varied entropy
        n_val = len(val_ground_truth)
        oracle_choices = val_ground_truth.argmax(axis=1)

        # Simulated accuracy: model is correct X% of the time
        model_accuracy = 0.35  # ~35% accuracy like we observed

        predictions = np.zeros_like(val_ground_truth)
        for i in range(n_val):
            # Decide if model will be correct on this query
            is_correct = np.random.random() < model_accuracy

            # Confidence varies by query
            confidence = 1.5 + np.random.exponential(1.0)  # Range ~1.5-5

            if is_correct:
                # Boost the oracle choice
                predictions[i, oracle_choices[i]] = confidence
                # Small random logits for others
                for j in range(n_llms):
                    if j != oracle_choices[i]:
                        predictions[i, j] = np.random.randn() * 0.3
            else:
                # Pick a wrong LLM with high confidence (mode collapse simulation)
                wrong_choices = [j for j in range(n_llms) if j != oracle_choices[i]]
                chosen_wrong = np.random.choice(wrong_choices)
                predictions[i, chosen_wrong] = confidence * 0.8
                for j in range(n_llms):
                    if j != chosen_wrong:
                        predictions[i, j] = np.random.randn() * 0.3

    # Run evaluation
    results = run_evaluation(
        predictions=predictions,
        ground_truth=val_ground_truth,
        llm_names=llm_names,
        costs=costs,
        latencies=latencies,
    )

    print_results(results)

    # === ANALYSIS BY ENTROPY BUCKET ===
    probs = F.softmax(torch.tensor(predictions), dim=1).numpy()
    entropy = compute_entropy(probs)

    print("\n" + "=" * 100)
    print("ENTROPY BUCKET ANALYSIS")
    print("=" * 100)

    q25, q50, q75 = np.percentile(entropy, [25, 50, 75])
    buckets = [
        ("Low (0-25%)", 0, q25),
        ("Medium-Low (25-50%)", q25, q50),
        ("Medium-High (50-75%)", q50, q75),
        ("High (75-100%)", q75, 1.01),
    ]

    print(f"\n{'Bucket':<25} {'N':>6} {'Entropy':>10} {'OracleAcc':>12} {'BestSingle':>12}")
    print("-" * 70)

    for label, lo, hi in buckets:
        mask = (entropy >= lo) & (entropy < hi)
        n = mask.sum()
        if n == 0:
            continue

        bucket_gt = val_ground_truth[mask]
        oracle_acc = bucket_gt.max(axis=1).mean()
        best_single = bucket_gt.mean(axis=0).max()
        avg_entropy = entropy[mask].mean()

        print(f"{label:<25} {n:>6} {avg_entropy:>10.3f} {oracle_acc*100:>11.1f}% {best_single*100:>11.1f}%")

    print("\nInterpretation:")
    print("  - High entropy queries are uncertain → routing has more value")
    print("  - Low entropy queries are confident → argmax should work well")


if __name__ == "__main__":
    main()
