#!/usr/bin/env python3
"""Simulate routing decisions and compute actual accuracy metrics.

This script answers: "If we route using the trained model, what actual
accuracy do we achieve on benchmarks like BBEH?"

Key insight: frac_oracle is a RELATIVE metric (% of oracle uplift captured).
This script computes ABSOLUTE accuracy by simulating routing decisions.
"""

import argparse
import logging
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import yaml

sys.path.insert(0, str(Path(__file__).parent.parent))

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)


def compute_routing_accuracy(
    predictions: np.ndarray,  # (num_queries, num_llms) logits
    ground_truth: np.ndarray,  # (num_queries, num_llms) binary success
    llm_names: list[str],
) -> dict:
    """Compute actual accuracy metrics from routing decisions.

    Args:
        predictions: Router logits (num_queries, num_llms)
        ground_truth: Binary success indicators (num_queries, num_llms)
        llm_names: Names of LLMs

    Returns:
        Dictionary with accuracy metrics
    """
    num_queries, num_llms = predictions.shape

    # Router's choices (argmax of predictions)
    router_choices = predictions.argmax(axis=1)

    # Oracle's choices (argmax of ground truth, tie-break by index)
    oracle_choices = ground_truth.argmax(axis=1)

    # Actual performance achieved by each strategy
    router_accuracy = np.array([
        ground_truth[i, router_choices[i]] for i in range(num_queries)
    ]).mean()

    oracle_accuracy = np.array([
        ground_truth[i, oracle_choices[i]] for i in range(num_queries)
    ]).mean()

    # Single-model baselines (always pick same LLM)
    single_model_accuracies = {}
    for llm_idx, llm_name in enumerate(llm_names):
        acc = ground_truth[:, llm_idx].mean()
        single_model_accuracies[llm_name] = acc

    best_single_model = max(single_model_accuracies, key=single_model_accuracies.get)
    best_single_accuracy = single_model_accuracies[best_single_model]

    # Random routing baseline
    random_accuracy = ground_truth.mean()

    # Worst single model
    worst_single_model = min(single_model_accuracies, key=single_model_accuracies.get)
    worst_single_accuracy = single_model_accuracies[worst_single_model]

    # Compute frac_oracle ourselves to verify
    # frac_oracle = mean((router_reward - min_reward) / (oracle_reward - min_reward))
    min_rewards = ground_truth.min(axis=1)
    router_rewards = np.array([ground_truth[i, router_choices[i]] for i in range(num_queries)])
    oracle_rewards = ground_truth.max(axis=1)

    # Handle ties correctly: for 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
        (router_rewards - min_rewards) / np.maximum(span, 1e-6),
    )
    frac_oracle = per_query_frac.mean()

    return {
        "router_accuracy": router_accuracy,
        "oracle_accuracy": oracle_accuracy,
        "best_single_model": best_single_model,
        "best_single_accuracy": best_single_accuracy,
        "worst_single_model": worst_single_model,
        "worst_single_accuracy": worst_single_accuracy,
        "random_accuracy": random_accuracy,
        "frac_oracle_computed": frac_oracle,
        "single_model_accuracies": single_model_accuracies,
        "num_queries": num_queries,
        "num_llms": num_llms,
        # Breakdown by query difficulty
        "router_choices": router_choices,
        "oracle_choices": oracle_choices,
    }


def analyze_by_task(
    predictions: np.ndarray,
    ground_truth: np.ndarray,
    task_ids: np.ndarray,
    llm_names: list[str],
) -> pd.DataFrame:
    """Analyze routing accuracy by task/benchmark."""
    unique_tasks = np.unique(task_ids)
    results = []

    for task in unique_tasks:
        mask = task_ids == task
        task_preds = predictions[mask]
        task_gt = ground_truth[mask]

        if len(task_preds) == 0:
            continue

        metrics = compute_routing_accuracy(task_preds, task_gt, llm_names)
        results.append({
            "task": task,
            "n_queries": metrics["num_queries"],
            "router_acc": metrics["router_accuracy"],
            "oracle_acc": metrics["oracle_accuracy"],
            "best_single_acc": metrics["best_single_accuracy"],
            "best_single_model": metrics["best_single_model"],
            "random_acc": metrics["random_accuracy"],
            "router_vs_best_single": metrics["router_accuracy"] - metrics["best_single_accuracy"],
            "router_vs_oracle": metrics["router_accuracy"] - metrics["oracle_accuracy"],
        })

    return pd.DataFrame(results).sort_values("router_vs_best_single", ascending=False)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="Config YAML path")
    parser.add_argument("--checkpoint", help="Model checkpoint (optional)")
    args = parser.parse_args()

    # Load config
    config_path = Path(args.config)
    with open(config_path) as f:
        config = yaml.safe_load(f)

    base_dir = config_path.parent.parent
    router_data_path = config.get("saved_router_data_path", "data/router_data.parquet")
    if not Path(router_data_path).is_absolute():
        router_data_path = str(base_dir / router_data_path)

    llm_num = config.get("llm_num", 7)

    logger.info("=" * 70)
    logger.info("ROUTING ACCURACY SIMULATION")
    logger.info("=" * 70)
    logger.info(f"Config: {args.config}")
    logger.info(f"Data: {router_data_path}")
    logger.info(f"LLMs: {llm_num}")

    # Load data
    logger.info("\nLoading data...")
    df = pd.read_parquet(router_data_path, columns=["query", "task_id", "llm", "effect"])

    # Get unique LLM names (first 7 only)
    llm_names = df["llm"].unique().tolist()[:llm_num]
    logger.info(f"LLMs: {llm_names}")

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

    # Group by (query, task_id) and pivot to get (query, llm) matrix
    df["group_key"] = df["query"].astype(str) + "|||" + df["task_id"].astype(str)

    # Only keep complete groups (exactly 7 LLMs)
    group_sizes = df.groupby("group_key").size()
    valid_groups = group_sizes[group_sizes == llm_num].index
    df = df[df["group_key"].isin(valid_groups)]

    logger.info(f"Valid query groups: {len(valid_groups)}")

    # Create ground truth matrix
    # First, create a mapping from LLM name to index
    llm_to_idx = {name: i for i, name in enumerate(llm_names)}

    # Group and extract
    ground_truth_list = []
    task_ids_list = []
    query_list = []

    for group_key, group_df in df.groupby("group_key", sort=False):
        gt_row = np.zeros(llm_num)
        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"]
        ground_truth_list.append(gt_row)
        task_ids_list.append(group_df["task_id"].iloc[0])
        query_list.append(group_df["query"].iloc[0])

    ground_truth = np.array(ground_truth_list)
    task_ids = np.array(task_ids_list)

    logger.info(f"Ground truth shape: {ground_truth.shape}")

    # === SIMULATE DIFFERENT ROUTING STRATEGIES ===

    # 1. Oracle (always best)
    oracle_predictions = ground_truth.copy()

    # 2. Random (uniform)
    np.random.seed(1881)
    random_predictions = np.random.rand(*ground_truth.shape)

    # 3. Always best single model (need to find which one)
    single_model_accs = ground_truth.mean(axis=0)
    best_single_idx = single_model_accs.argmax()
    best_single_predictions = np.zeros_like(ground_truth)
    best_single_predictions[:, best_single_idx] = 1.0

    # 4. Simulated router at different frac_oracle levels
    # Key insight: if frac_oracle = X, then on average:
    # router_reward = min_reward + X * (oracle_reward - min_reward)

    logger.info("\n" + "=" * 70)
    logger.info("STRATEGY COMPARISON (Absolute Accuracy)")
    logger.info("=" * 70)

    # Oracle analysis
    oracle_metrics = compute_routing_accuracy(oracle_predictions, ground_truth, llm_names)
    logger.info(f"\n{'Oracle (always best):':<30} {oracle_metrics['oracle_accuracy']*100:.2f}%")

    # Best single model
    logger.info(f"{'Best single (' + llm_names[best_single_idx] + '):':<30} {single_model_accs[best_single_idx]*100:.2f}%")

    # All single models
    logger.info("\nSingle model accuracies:")
    for i, name in enumerate(llm_names):
        logger.info(f"  {name:<25} {single_model_accs[i]*100:.2f}%")

    # Random
    random_metrics = compute_routing_accuracy(random_predictions, ground_truth, llm_names)
    logger.info(f"\n{'Random routing:':<30} {random_metrics['router_accuracy']*100:.2f}%")

    # === KEY CALCULATION ===
    # What accuracy does frac_oracle = 72.4% translate to?

    logger.info("\n" + "=" * 70)
    logger.info("TRANSLATING frac_oracle TO ABSOLUTE ACCURACY")
    logger.info("=" * 70)

    # Per-query: router_reward = min + frac_oracle * (max - min)
    min_rewards = ground_truth.min(axis=1)
    max_rewards = ground_truth.max(axis=1)
    span = max_rewards - min_rewards

    # For queries where all LLMs agree (span=0), routing doesn't matter
    informative_queries = span > 0
    logger.info(f"\nQueries where routing matters: {informative_queries.sum()} / {len(span)} ({100*informative_queries.mean():.1f}%)")
    logger.info(f"Queries where all LLMs same: {(~informative_queries).sum()} ({100*(~informative_queries).mean():.1f}%)")

    # Expected accuracy at different frac_oracle levels
    logger.info("\nExpected accuracy at different frac_oracle levels:")
    logger.info("-" * 50)

    for fo in [0.5, 0.6, 0.7, 0.724, 0.8, 0.9, 1.0]:
        # Expected reward = min + fo * (max - min)
        expected_reward = min_rewards + fo * span
        expected_accuracy = expected_reward.mean()
        vs_best_single = expected_accuracy - single_model_accs[best_single_idx]
        logger.info(f"  frac_oracle={fo:.1%}: accuracy={expected_accuracy*100:.2f}% (vs best single: {vs_best_single*100:+.2f}%)")

    # === BREAKDOWN BY TASK ===
    logger.info("\n" + "=" * 70)
    logger.info("ACCURACY BY TASK/BENCHMARK")
    logger.info("=" * 70)

    # Group tasks
    task_stats = []
    for task in np.unique(task_ids):
        mask = task_ids == task
        task_gt = ground_truth[mask]
        n = mask.sum()
        oracle_acc = task_gt.max(axis=1).mean()
        best_single = task_gt.mean(axis=0).max()
        random_acc = task_gt.mean()
        span_mean = (task_gt.max(axis=1) - task_gt.min(axis=1)).mean()

        task_stats.append({
            "task": task,
            "n": n,
            "oracle": oracle_acc,
            "best_single": best_single,
            "random": random_acc,
            "span": span_mean,
        })

    task_df = pd.DataFrame(task_stats).sort_values("span", ascending=False)

    logger.info("\nTasks sorted by routing opportunity (span):")
    logger.info(f"{'Task':<20} {'N':>6} {'Oracle':>8} {'BestSgl':>8} {'Random':>8} {'Span':>8}")
    logger.info("-" * 60)
    for _, row in task_df.head(15).iterrows():
        logger.info(f"{row['task'][:20]:<20} {row['n']:>6} {row['oracle']*100:>7.1f}% {row['best_single']*100:>7.1f}% {row['random']*100:>7.1f}% {row['span']:>7.2f}")

    # === FINAL SUMMARY ===
    logger.info("\n" + "=" * 70)
    logger.info("SUMMARY: What 72.4% frac_oracle means")
    logger.info("=" * 70)

    fo = 0.724
    expected_reward = min_rewards + fo * span
    router_accuracy = expected_reward.mean()

    logger.info(f"\nOn this dataset ({len(ground_truth)} queries, {llm_num} LLMs):")
    logger.info(f"  Oracle accuracy:      {oracle_metrics['oracle_accuracy']*100:.2f}%")
    logger.info(f"  Router accuracy:      {router_accuracy*100:.2f}% (at frac_oracle=72.4%)")
    logger.info(f"  Best single model:    {single_model_accs[best_single_idx]*100:.2f}% ({llm_names[best_single_idx]})")
    logger.info(f"  Random routing:       {ground_truth.mean()*100:.2f}%")

    logger.info(f"\nRouter improvement over best single model:")
    improvement = router_accuracy - single_model_accs[best_single_idx]
    logger.info(f"  Absolute: {improvement*100:+.2f}%")
    logger.info(f"  Relative: {100*improvement/single_model_accs[best_single_idx]:+.1f}% relative improvement")

    # Gap analysis
    gap_to_oracle = oracle_metrics['oracle_accuracy'] - router_accuracy
    gap_closed = router_accuracy - single_model_accs[best_single_idx]
    total_gap = oracle_metrics['oracle_accuracy'] - single_model_accs[best_single_idx]

    logger.info(f"\nGap analysis:")
    logger.info(f"  Total gap (oracle - best single):     {total_gap*100:.2f}%")
    logger.info(f"  Gap closed by router:                 {gap_closed*100:.2f}% ({100*gap_closed/total_gap:.1f}% of gap)")
    logger.info(f"  Remaining gap (oracle - router):      {gap_to_oracle*100:.2f}%")


if __name__ == "__main__":
    main()
