#!/usr/bin/env python3
"""Deep error analysis for needle-weighted LLM router.

This script analyzes the performance gap between the needle-weighted model
and the oracle, using the actual continuous reward structure.

Key questions addressed:
1. What is the distribution of oracle rewards (best possible)?
2. How does model performance vary by ground-truth difficulty?
3. What specific patterns cause the largest regret?
"""

import pandas as pd
import numpy as np
from collections import defaultdict

# Load the router data
print("Loading router_data.csv...")
df = pd.read_csv("data/router_data.csv")

print(f"Total rows: {len(df):,}")

# Pivot to get LLM performance per query
print("\n" + "="*70)
print("GROUND TRUTH DIFFICULTY ANALYSIS")
print("="*70)

query_cols = ["task_id", "query"]

# Get rewards per LLM for each query
query_llm_perf = df.pivot_table(
    index=["task_id", "query"],
    columns="llm",
    values="effect",
    aggfunc="first"
).reset_index()

llm_cols = [c for c in query_llm_perf.columns if c not in ["task_id", "query"]]
print(f"Unique queries: {len(query_llm_perf):,}")
print(f"LLMs: {llm_cols}")

# For each query, compute key statistics
query_llm_perf["oracle_reward"] = query_llm_perf[llm_cols].max(axis=1)
query_llm_perf["min_reward"] = query_llm_perf[llm_cols].min(axis=1)
query_llm_perf["mean_reward"] = query_llm_perf[llm_cols].mean(axis=1)
query_llm_perf["reward_range"] = query_llm_perf["oracle_reward"] - query_llm_perf["min_reward"]

# Ground-truth difficulty: how spread out are the rewards?
# Low spread = easy (all LLMs similar), high spread = routing matters
query_llm_perf["reward_std"] = query_llm_perf[llm_cols].std(axis=1)

# Define "needle" as queries where ONE LLM is clearly better
# We'll use the gap between best and 2nd best
sorted_rewards = np.sort(query_llm_perf[llm_cols].values, axis=1)[:, ::-1]  # descending
query_llm_perf["top1_reward"] = sorted_rewards[:, 0]
query_llm_perf["top2_reward"] = sorted_rewards[:, 1]
query_llm_perf["top1_top2_gap"] = query_llm_perf["top1_reward"] - query_llm_perf["top2_reward"]

n_total = len(query_llm_perf)

print("\n📊 Oracle Reward Distribution:")
print("-" * 70)
oracle_stats = query_llm_perf["oracle_reward"].describe()
print(f"  Mean oracle reward:   {oracle_stats['mean']:.4f}")
print(f"  Std:                  {oracle_stats['std']:.4f}")
print(f"  Min:                  {oracle_stats['min']:.4f}")
print(f"  25%:                  {oracle_stats['25%']:.4f}")
print(f"  50% (median):         {oracle_stats['50%']:.4f}")
print(f"  75%:                  {oracle_stats['75%']:.4f}")
print(f"  Max:                  {oracle_stats['max']:.4f}")

# Key: What fraction of queries have oracle_reward = 0 (no LLM succeeds at all)?
zero_oracle_mask = query_llm_perf["oracle_reward"] == 0
n_zero_oracle = zero_oracle_mask.sum()
print(f"\n  Queries with oracle_reward=0: {n_zero_oracle:,} ({n_zero_oracle/n_total*100:.1f}%)")

# What fraction have low oracle (< 0.5)?
low_oracle_mask = query_llm_perf["oracle_reward"] < 0.5
n_low_oracle = low_oracle_mask.sum()
print(f"  Queries with oracle_reward<0.5: {n_low_oracle:,} ({n_low_oracle/n_total*100:.1f}%)")

# Distribution of top1-top2 gap (routing impact)
print("\n📊 Routing Impact Distribution (top1-top2 gap):")
print("-" * 70)
gap_stats = query_llm_perf["top1_top2_gap"].describe()
print(f"  Mean gap:   {gap_stats['mean']:.4f}")
print(f"  Std:        {gap_stats['std']:.4f}")
print(f"  Min:        {gap_stats['min']:.4f}")
print(f"  Median:     {gap_stats['50%']:.4f}")
print(f"  Max:        {gap_stats['max']:.4f}")

# Define needle queries: large gap between #1 and #2
NEEDLE_THRESHOLD = 0.3  # Gap > 30% means clear winner
needle_mask = query_llm_perf["top1_top2_gap"] > NEEDLE_THRESHOLD
n_needle = needle_mask.sum()
print(f"\n  'Needle' queries (gap > {NEEDLE_THRESHOLD}): {n_needle:,} ({n_needle/n_total*100:.1f}%)")

# Trivial queries: all LLMs very similar
TRIVIAL_THRESHOLD = 0.05  # Gap < 5% means doesn't matter
trivial_mask = query_llm_perf["top1_top2_gap"] < TRIVIAL_THRESHOLD
n_trivial = trivial_mask.sum()
print(f"  'Trivial' queries (gap < {TRIVIAL_THRESHOLD}): {n_trivial:,} ({n_trivial/n_total*100:.1f}%)")

# Calculate theoretical ceiling and gaps
print("\n" + "="*70)
print("THEORETICAL CEILING AND MODEL PERFORMANCE")
print("="*70)

# Oracle = always pick best LLM
oracle_total = query_llm_perf["oracle_reward"].mean()
print(f"\nOracle (always best): {oracle_total:.4f}")

# Random = uniform random choice
random_total = query_llm_perf["mean_reward"].mean()
print(f"Random (uniform):     {random_total:.4f}")

# Worst = always pick worst
worst_total = query_llm_perf["min_reward"].mean()
print(f"Worst (always worst): {worst_total:.4f}")

# Model achieved 69.14% frac_of_oracle
# frac_of_oracle = (chosen - min) / (oracle - min)
# So: chosen = min + frac_of_oracle * (oracle - min)
NEEDLE_MODEL_FRAC = 0.6914

# Reconstruct average chosen reward
range_total = query_llm_perf["reward_range"].mean()
model_chosen_approx = worst_total + NEEDLE_MODEL_FRAC * range_total
print(f"\nNeedle model (69.14%): ~{model_chosen_approx:.4f} average reward")

# Gap analysis
print(f"\nGap to oracle: {oracle_total - model_chosen_approx:.4f} reward units")
print(f"Frac of max gap: {(oracle_total - model_chosen_approx) / (oracle_total - random_total) * 100:.1f}%")

# Stratify by ground-truth difficulty
print("\n" + "="*70)
print("PERFORMANCE BY GROUND-TRUTH DIFFICULTY")
print("="*70)

# Define difficulty quartiles based on top1-top2 gap
# High gap = easy to route (clear winner)
# Low gap = hard to route (all similar)
query_llm_perf["routing_difficulty"] = pd.qcut(
    query_llm_perf["top1_top2_gap"].rank(method='first'),
    q=4,
    labels=["Hard (Q4)", "Medium-Hard (Q3)", "Medium-Easy (Q2)", "Easy (Q1)"]
)

print("\n📊 Performance by GROUND-TRUTH difficulty (not model confidence!):")
print("-" * 70)
for diff_level in ["Easy (Q1)", "Medium-Easy (Q2)", "Medium-Hard (Q3)", "Hard (Q4)"]:
    mask = query_llm_perf["routing_difficulty"] == diff_level
    subset = query_llm_perf[mask]

    n = len(subset)
    oracle_r = subset["oracle_reward"].mean()
    random_r = subset["mean_reward"].mean()
    min_r = subset["min_reward"].mean()
    avg_gap = subset["top1_top2_gap"].mean()

    # Theoretical model performance at 69.14% frac_oracle
    range_r = subset["reward_range"].mean()
    model_r = min_r + NEEDLE_MODEL_FRAC * range_r

    print(f"\n{diff_level}:")
    print(f"  Queries:     {n:>6,} ({n/n_total*100:>5.1f}%)")
    print(f"  Avg gap:     {avg_gap:>6.3f}")
    print(f"  Oracle:      {oracle_r:>6.3f}")
    print(f"  Random:      {random_r:>6.3f}")
    print(f"  Model@69%:   {model_r:>6.3f}")
    print(f"  Lift over random: {(model_r - random_r)/(oracle_r - random_r + 1e-8)*100:>5.1f}%")

# Which LLM is best for which task?
print("\n" + "="*70)
print("LLM WIN RATES")
print("="*70)

print("\n📊 Which LLM is oracle most often?")
print("-" * 70)

# Find the winning LLM for each query
def get_winner(row):
    llm_rewards = {llm: row[llm] for llm in llm_cols}
    return max(llm_rewards, key=llm_rewards.get)

query_llm_perf["winner"] = query_llm_perf.apply(get_winner, axis=1)

winner_counts = query_llm_perf["winner"].value_counts()
for llm, count in winner_counts.items():
    pct = count / n_total * 100
    bar = "█" * int(pct / 2)
    print(f"  {llm:<30} {count:>5,} ({pct:>5.1f}%) {bar}")

# LLM win rates on needle queries specifically
print("\n📊 LLM wins on NEEDLE queries (gap > 0.3):")
print("-" * 70)

needle_winners = query_llm_perf[needle_mask]["winner"].value_counts()
for llm, count in needle_winners.items():
    pct = count / n_needle * 100
    bar = "█" * int(pct / 2)
    print(f"  {llm:<30} {count:>5,} ({pct:>5.1f}%) {bar}")

# Per-task analysis
print("\n" + "="*70)
print("PER-TASK ANALYSIS")
print("="*70)

task_stats = []
for task_id in query_llm_perf["task_id"].unique():
    task_mask = query_llm_perf["task_id"] == task_id
    task_df = query_llm_perf[task_mask]

    task_stats.append({
        "task_id": task_id,
        "n_queries": len(task_df),
        "avg_oracle": task_df["oracle_reward"].mean(),
        "avg_gap": task_df["top1_top2_gap"].mean(),
        "needle_pct": (task_df["top1_top2_gap"] > NEEDLE_THRESHOLD).sum() / len(task_df) * 100,
        "trivial_pct": (task_df["top1_top2_gap"] < TRIVIAL_THRESHOLD).sum() / len(task_df) * 100,
        "dominant_llm": task_df["winner"].mode().iloc[0] if len(task_df) > 0 else "N/A",
    })

task_stats_df = pd.DataFrame(task_stats)

# Sort by needle percentage (hardest routing problems)
task_stats_df = task_stats_df.sort_values("needle_pct", ascending=False)

print("\nTasks ranked by NEEDLE query percentage (hardest for routing):")
print("-" * 70)
for _, row in task_stats_df.head(15).iterrows():
    print(f"  {row['task_id']:<35} | needle: {row['needle_pct']:>5.1f}% | "
          f"oracle: {row['avg_oracle']:>.3f} | best: {row['dominant_llm'][:15]}")

# KEY INSIGHT: What's the theoretical max with perfect calibration?
print("\n" + "="*70)
print("KEY INSIGHTS")
print("="*70)

print(f"""
1. ORACLE PERFORMANCE: {oracle_total:.3f}
   - This is the theoretical ceiling if we always pick the best LLM

2. RANDOM BASELINE: {random_total:.3f}
   - Picking uniformly at random would achieve this

3. NEEDLE MODEL: ~{model_chosen_approx:.3f} (69.14% of normalized gap)
   - Captures {(model_chosen_approx - random_total)/(oracle_total - random_total)*100:.1f}% of the oracle-random gap

4. WHERE ROUTING MATTERS:
   - Needle queries ({n_needle:,} = {n_needle/n_total*100:.1f}%): Clear winner, routing critical
   - Trivial queries ({n_trivial:,} = {n_trivial/n_total*100:.1f}%): All LLMs similar, routing doesn't matter

5. IMPROVEMENT OPPORTUNITY:
   - Focus on the {n_needle:,} needle queries
   - Gap of {oracle_total - model_chosen_approx:.3f} reward units to recover
   - Primarily affects: {', '.join(task_stats_df.head(5)['task_id'].tolist())}
""")

# Save results
import json
results = {
    "total_queries": int(n_total),
    "oracle_avg": float(oracle_total),
    "random_avg": float(random_total),
    "model_frac_oracle": float(NEEDLE_MODEL_FRAC),
    "model_reward_approx": float(model_chosen_approx),
    "needle_queries": int(n_needle),
    "needle_pct": float(n_needle / n_total),
    "trivial_queries": int(n_trivial),
    "trivial_pct": float(n_trivial / n_total),
    "gap_to_oracle": float(oracle_total - model_chosen_approx),
    "lift_over_random_pct": float((model_chosen_approx - random_total) / (oracle_total - random_total) * 100),
}

with open("results/error_analysis_deep.json", "w") as f:
    json.dump(results, f, indent=2)

print("\n✅ Results saved to results/error_analysis_deep.json")
