#!/usr/bin/env python3
"""Comprehensive difficulty-aware evaluation for GNN router.

This script evaluates router performance stratified by query difficulty,
validates uncertainty calibration, and confirms that GNN predictions
correlate with ground truth difficulty patterns.

Key metrics computed:
1. Difficulty Distribution: Queries grouped by how many models got them correct
2. Difficulty-Stratified Metrics: Top-1 acc, regret, NDCG per difficulty band
3. Uncertainty Calibration: ECE, entropy vs difficulty correlation
4. GNN Difficulty Prediction: Correlation with ground truth easiness
5. Ranking Discrimination: Spearman correlation, per-LLM/task affinity
6. Advanced Analysis (--advanced): Error categorization, cost-quality tradeoffs,
   statistical significance via bootstrap

Usage:
    python scripts/eval_difficulty_aware.py --checkpoint registry/eb1-gnn-preview-12-19
    python scripts/eval_difficulty_aware.py --checkpoint registry/model --advanced
    python scripts/eval_difficulty_aware.py --data-only  # Ground truth analysis only

Note on Inference Benchmarks:
    Use --run-benchmark on the target production hardware (e.g., B200 GPU instance)
    to get accurate latency/throughput metrics. Local benchmarks are not representative
    of production performance.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# Ensure embergraphrouter is importable when running as a script
# Add eb1 directory (parent of embergraphrouter) to path so 'embergraphrouter' is a package
_eb1_dir = Path(__file__).resolve().parent.parent.parent
if str(_eb1_dir) not in sys.path:
    sys.path.insert(0, str(_eb1_dir))

import numpy as np
import pandas as pd
import torch

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


def _json_default(obj: object) -> object:
    """Convert numpy/torch types to JSON-serializable Python types."""
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, (np.bool_,)):
        return bool(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


# Difficulty bands based on number of models correct (out of N LLMs)
DIFFICULTY_BANDS = {
    "impossible": (0, 0),      # No model correct
    "very_hard": (1, 2),       # 1-2 models correct
    "hard": (3, 4),            # 3-4 models correct
    "medium": (5, 6),          # 5-6 models correct
    "easy": (7, 7),            # 7 models correct
    "trivial": (8, float("inf")),  # All models correct
}


@dataclass
class QueryDifficultyStats:
    """Statistics for a single query's difficulty."""

    query_id: int
    query_text: str
    task_id: str
    n_correct: int
    pct_correct: float
    effect_mean: float
    effect_std: float
    effect_min: float
    effect_max: float
    perf_range: float  # max - min effect (routing value)
    agreement_entropy: float  # Binary entropy of correct/incorrect
    difficulty_band: str


@dataclass
class DifficultyBandStats:
    """Aggregate statistics for a difficulty band."""

    band_name: str
    n_queries: int
    pct_of_total: float
    avg_effect: float
    avg_effect_std: float
    avg_perf_range: float
    avg_entropy: float


@dataclass
class RoutingMetrics:
    """Routing performance metrics."""

    top1_acc: float
    regret: float
    frac_of_oracle: float
    ndcg_at_1: float
    ndcg_at_3: float
    avg_chosen_reward: float
    avg_oracle_reward: float


@dataclass
class CalibrationMetrics:
    """Uncertainty calibration metrics.

    These metrics test whether the model "knows when it doesn't know":
    - entropy_difficulty_corr: Does entropy reflect query difficulty?
    - entropy_error_corr: Does high entropy predict routing mistakes?
    - entropy_quality_corr: Does low entropy predict high routing quality?

    A well-calibrated router should have:
    - Positive entropy_error_corr (high entropy -> more mistakes)
    - Negative entropy_quality_corr (high entropy -> lower frac_of_oracle)

    Raw Logit Uncertainty Metrics:
    Since softmax compresses small logit differences into near-uniform probabilities,
    we also compute uncertainty directly from raw logits:
    - logit_std: Standard deviation of logits (high = confident, low = uncertain)
    - logit_margin: max_logit - second_max (high = confident in top choice)
    """

    ece: float  # Expected Calibration Error
    mce: float  # Maximum Calibration Error
    overconfidence_rate: float  # P(wrong | high_confidence)
    underconfidence_rate: float  # P(correct | low_confidence)
    entropy_difficulty_corr: float  # Correlation(entropy, num_correct)
    confidence_difficulty_corr: float  # Correlation(max_prob, num_correct)
    # "Knows when it doesn't know" metrics (softmax-based entropy)
    entropy_error_corr: float  # Correlation(entropy, routing_error) - should be POSITIVE
    entropy_quality_corr: float  # Correlation(entropy, frac_of_oracle) - should be NEGATIVE
    avg_entropy: float  # Average entropy (near max=log(n_llms) means near-uniform)
    entropy_std: float  # Std of entropy (low std means no differentiation)
    pct_high_entropy: float  # % queries with entropy > 0.9 * max_entropy
    # Raw logit uncertainty metrics (bypass softmax compression)
    logit_std_error_corr: float  # Corr(logit_std, error) - should be NEGATIVE (high std = confident)
    logit_std_quality_corr: float  # Corr(logit_std, quality) - should be POSITIVE
    logit_margin_error_corr: float  # Corr(margin, error) - should be NEGATIVE (high margin = confident)
    logit_margin_quality_corr: float  # Corr(margin, quality) - should be POSITIVE
    avg_logit_std: float  # Average std of logits per query
    avg_logit_margin: float  # Average margin (max - 2nd max) per query


@dataclass
class EdgeCaseCalibration:
    """Calibration metrics for impossible and trivial queries.

    A well-calibrated model should:
    - Show HIGH entropy on impossible queries (all models fail, no differentiation)
    - Show LOW entropy on trivial queries (all models succeed, no differentiation)
    - NOT be overconfident on impossible queries (predicting success when all fail)
    - Route impossible/trivial to cheap models (cost optimization opportunity)
    """

    # Impossible queries (0/N correct)
    n_impossible: int
    pct_impossible: float
    impossible_avg_entropy: float  # Should be HIGH (near max_entropy)
    impossible_avg_max_prob: float  # Should be LOW (uncertain)
    impossible_overconfidence: float  # P(max_prob > 0.7) on impossible - should be LOW

    # Trivial queries (N/N correct)
    n_trivial: int
    pct_trivial: float
    trivial_avg_entropy: float  # Should be LOW (all models good)
    trivial_avg_max_prob: float  # Can be high (confident, and correct)
    trivial_routing_variance: float  # Variance in predictions - should be LOW

    # Cost optimization metrics
    cost_opt_safe_cheap_pct: float  # % of impossible+trivial safely routable to cheap
    cost_opt_potential_savings: float  # Estimated cost savings if routing edge cases to cheap


@dataclass
class DifficultyDetection:
    """Difficulty detection using mean logit signal.

    Key insight: While entropy is near-uniform for all queries (model outputs
    near-uniform probabilities), the RAW MEAN LOGIT captures difficulty:
    - High mean logit (+0.5 to +1.0): Trivial queries (all LLMs likely correct)
    - Near-zero mean logit: Medium/hard queries (mixed LLM performance)
    - Negative mean logit (-0.3 to -0.6): Very hard/impossible queries

    This enables a production routing policy:
    - mean_logit > threshold_high: Route to cheapest (trivial, any LLM works)
    - mean_logit < threshold_low: Route to best-on-hard LLMs or cheapest
    - Otherwise: Use GNN's top-1 prediction

    The mean logit represents the model's belief about P(any LLM correct),
    which directly measures query difficulty.
    """

    # Mean logit statistics by difficulty band
    mean_logit_by_band: dict[str, float]  # band -> avg mean_logit
    mean_logit_std_by_band: dict[str, float]  # band -> std of mean_logit

    # Separation quality for trivial vs impossible
    trivial_vs_impossible_auc: float  # AUC for separating trivial from impossible
    optimal_threshold: float  # Best threshold for separation
    separation_accuracy: float  # Accuracy at optimal threshold

    # Correlation with ground truth difficulty
    mean_logit_difficulty_corr: float  # Correlation(mean_logit, n_correct)

    # Recommended thresholds for production
    recommended_trivial_threshold: float  # mean_logit > this -> likely trivial
    recommended_hard_threshold: float  # mean_logit < this -> likely very_hard/impossible

    # Best LLMs for hard queries (for fallback routing)
    best_llms_on_hard: list[str]  # Top LLMs by success rate on very_hard queries
    best_llm_success_rates: dict[str, float]  # LLM -> success rate on very_hard

    def get_routing_decision(
        self,
        mean_logit: float,
        gnn_top1_llm: str,
        cheap_llm: str = "GPT-5 Nano",
    ) -> tuple[str, str]:
        """Get production routing decision based on mean logit.

        This implements the recommended routing policy:
        1. If mean_logit > trivial_threshold: Route to cheapest (cost savings)
        2. If mean_logit < hard_threshold: Route to best-on-hard LLM (quality)
        3. Otherwise: Use GNN's top-1 prediction (standard routing)

        Args:
            mean_logit: Mean logit from model prediction for this query
            gnn_top1_llm: The LLM that GNN predicts as best
            cheap_llm: Name of cheapest LLM for cost optimization

        Returns:
            Tuple of (selected_llm, reason)
        """
        if mean_logit > self.recommended_trivial_threshold:
            return (cheap_llm, "trivial_detected")
        elif mean_logit < self.recommended_hard_threshold:
            fallback = self.best_llms_on_hard[0] if self.best_llms_on_hard else gnn_top1_llm
            return (fallback, "hard_detected")
        else:
            return (gnn_top1_llm, "gnn_routing")


@dataclass
class RankingDiscrimination:
    """Metrics for whether the router's rankings are discriminative.

    These answer: "Are the router's predictions meaningful?"

    Key questions:
    1. Does high p(LLM) mean that LLM performs better?
    2. Does the router capture task-specific LLM affinities?
    3. Are rankings correlated with ground truth per-query?

    A discriminative router should have:
    - High ranking_correlation (Spearman corr between predicted and true ranks)
    - Positive per_llm_discrimination (high score → better performance)
    - High task_affinity_captured (router learns which LLMs excel on which tasks)
    """

    # Per-query ranking correlation
    avg_spearman_corr: float  # Avg Spearman correlation per query (pred vs true ranks)
    avg_kendall_tau: float  # Avg Kendall tau per query
    pct_positive_corr: float  # % of queries with positive rank correlation

    # Per-LLM discrimination: when router gives high score, is LLM actually better?
    per_llm_discrimination: dict[str, float]  # LLM -> corr(router_score, actual_reward)
    avg_llm_discrimination: float  # Average across LLMs

    # Task-specific affinity capture
    task_affinity_captured: dict[str, float]  # task -> corr(router_rank, true_rank)
    avg_task_affinity: float  # Average across tasks
    tasks_with_strong_affinity: list[str]  # Tasks where router captures affinity well

    # Score-to-reward calibration (not just top-1)
    score_reward_corr: float  # Overall corr(all_scores, all_rewards)
    high_score_success_rate: float  # P(correct | score in top 25%)
    low_score_failure_rate: float  # P(incorrect | score in bottom 25%)


@dataclass
class ErrorAnalysis:
    """Detailed analysis of routing errors and failure modes.

    Provides insight into where and why the router fails, including
    error categorization by difficulty band, confusion patterns between
    LLMs, and identification of worst-case routing decisions.
    """

    # Overall error statistics
    n_errors: int  # Number of queries where router picked wrong LLM
    error_rate: float  # n_errors / n_queries
    avg_error_regret: float  # Average regret on error cases

    # Error categorization by difficulty
    errors_by_band: dict[str, int]  # band -> count of errors
    error_rate_by_band: dict[str, float]  # band -> error rate

    # Worst errors (highest regret)
    worst_error_regrets: list[float]  # Top 10 highest regret values
    worst_error_indices: list[int]  # Query indices of worst errors

    # Confusion patterns: which LLM did we pick vs should have picked?
    confusion_matrix: dict[str, dict[str, int]]  # predicted -> actual -> count
    most_confused_pairs: list[tuple[str, str, int]]  # (pred, actual, count) top 5

    # Error concentration: are errors spread out or concentrated?
    error_gini: float  # Gini coefficient of errors across tasks (0=uniform, 1=concentrated)
    tasks_with_most_errors: list[tuple[str, int]]  # (task, count) top 5


@dataclass
class CostQualityAnalysis:
    """Cost-quality tradeoff analysis for LLM routing.

    Measures the dollar impact of routing decisions by comparing
    router cost against baselines (random, oracle, cheapest-always)
    and evaluating quality at different cost budgets.
    """

    # LLM cost model (relative costs, normalized to cheapest=1.0)
    llm_costs: dict[str, float]  # LLM -> relative cost

    # Baseline comparisons
    random_routing_cost: float  # Expected cost if routing randomly
    oracle_routing_cost: float  # Cost if always picking best LLM
    router_routing_cost: float  # Actual cost with this router
    cheapest_always_cost: float  # Cost if always picking cheapest

    # Cost savings
    cost_savings_vs_random: float  # (random - router) / random
    cost_savings_vs_oracle: float  # How close to oracle? (0 = same as oracle)

    # Quality at different cost budgets
    quality_at_budget: dict[str, float]  # budget_level -> frac_of_oracle

    # Pareto efficiency
    is_pareto_optimal: bool  # Is this router on the Pareto frontier?
    pareto_gap: float  # Distance to Pareto frontier (0 = on frontier)


@dataclass
class StatisticalSignificance:
    """Statistical significance and confidence intervals for key metrics.

    Provides bootstrap confidence intervals and standard errors to
    determine whether observed differences are statistically meaningful.
    """

    # Confidence intervals (95%) for key metrics
    frac_of_oracle_ci: tuple[float, float]  # (lower, upper)
    regret_ci: tuple[float, float]
    spearman_ci: tuple[float, float]

    # Bootstrap standard errors
    frac_of_oracle_se: float
    regret_se: float
    spearman_se: float

    # Sample size and effective sample size
    n_samples: int
    n_effective: float  # Accounts for correlation


@dataclass
class InferenceBenchmark:
    """Inference performance benchmarks for production readiness.

    Measures latency percentiles, throughput, memory usage, and
    device characteristics to assess production viability.
    """

    # Latency (in milliseconds)
    latency_mean_ms: float
    latency_p50_ms: float
    latency_p95_ms: float
    latency_p99_ms: float

    # Throughput
    queries_per_second: float
    batch_size_tested: int

    # Memory
    model_size_mb: float
    peak_memory_mb: float

    # Device info
    device: str  # "cpu", "cuda", etc.
    dtype: str  # "float32", "float16", etc.


@dataclass
class EvaluationResults:
    """Complete evaluation results."""

    # Ground truth analysis
    n_queries: int
    n_llms: int
    difficulty_distribution: dict[str, DifficultyBandStats]
    task_difficulty: dict[str, float]  # task_id -> avg difficulty

    # Routing metrics (overall and per-band)
    overall_routing: RoutingMetrics | None = None
    routing_by_band: dict[str, RoutingMetrics] = field(default_factory=dict)

    # Calibration metrics
    calibration: CalibrationMetrics | None = None
    edge_case_calibration: EdgeCaseCalibration | None = None

    # GNN difficulty prediction
    difficulty_pred_corr: float | None = None  # Correlation with ground truth
    uncertainty_error_corr: float | None = None  # Correlation with prediction error

    # Difficulty detection using mean logit (key insight from analysis)
    difficulty_detection: DifficultyDetection | None = None

    # Ranking discrimination: are the router's rankings meaningful?
    ranking_discrimination: RankingDiscrimination | None = None

    # Advanced analysis
    error_analysis: ErrorAnalysis | None = None
    cost_quality: CostQualityAnalysis | None = None
    statistical_significance: StatisticalSignificance | None = None
    inference_benchmark: InferenceBenchmark | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert to JSON-serializable dict."""
        result = {
            "n_queries": self.n_queries,
            "n_llms": self.n_llms,
            "difficulty_distribution": {
                k: {
                    "n_queries": v.n_queries,
                    "pct_of_total": v.pct_of_total,
                    "avg_effect": v.avg_effect,
                    "avg_effect_std": v.avg_effect_std,
                    "avg_perf_range": v.avg_perf_range,
                    "avg_entropy": v.avg_entropy,
                }
                for k, v in self.difficulty_distribution.items()
            },
            "task_difficulty": self.task_difficulty,
        }

        if self.overall_routing:
            result["overall_routing"] = {
                "top1_acc": self.overall_routing.top1_acc,
                "regret": self.overall_routing.regret,
                "frac_of_oracle": self.overall_routing.frac_of_oracle,
                "ndcg_at_1": self.overall_routing.ndcg_at_1,
                "ndcg_at_3": self.overall_routing.ndcg_at_3,
            }

        if self.routing_by_band:
            result["routing_by_band"] = {
                k: {
                    "top1_acc": v.top1_acc,
                    "regret": v.regret,
                    "frac_of_oracle": v.frac_of_oracle,
                }
                for k, v in self.routing_by_band.items()
            }

        if self.calibration:
            result["calibration"] = {
                "ece": self.calibration.ece,
                "mce": self.calibration.mce,
                "overconfidence_rate": self.calibration.overconfidence_rate,
                "underconfidence_rate": self.calibration.underconfidence_rate,
                "entropy_difficulty_corr": self.calibration.entropy_difficulty_corr,
                "confidence_difficulty_corr": self.calibration.confidence_difficulty_corr,
                # "Knows when it doesn't know" metrics (softmax entropy)
                "entropy_error_corr": self.calibration.entropy_error_corr,
                "entropy_quality_corr": self.calibration.entropy_quality_corr,
                "avg_entropy": self.calibration.avg_entropy,
                "entropy_std": self.calibration.entropy_std,
                "pct_high_entropy": self.calibration.pct_high_entropy,
                # Raw logit uncertainty (bypass softmax)
                "logit_std_error_corr": self.calibration.logit_std_error_corr,
                "logit_std_quality_corr": self.calibration.logit_std_quality_corr,
                "logit_margin_error_corr": self.calibration.logit_margin_error_corr,
                "logit_margin_quality_corr": self.calibration.logit_margin_quality_corr,
                "avg_logit_std": self.calibration.avg_logit_std,
                "avg_logit_margin": self.calibration.avg_logit_margin,
            }

        if self.edge_case_calibration:
            ecc = self.edge_case_calibration
            result["edge_case_calibration"] = {
                "impossible": {
                    "n_queries": ecc.n_impossible,
                    "pct_of_total": ecc.pct_impossible,
                    "avg_entropy": ecc.impossible_avg_entropy,
                    "avg_max_prob": ecc.impossible_avg_max_prob,
                    "overconfidence_rate": ecc.impossible_overconfidence,
                },
                "trivial": {
                    "n_queries": ecc.n_trivial,
                    "pct_of_total": ecc.pct_trivial,
                    "avg_entropy": ecc.trivial_avg_entropy,
                    "avg_max_prob": ecc.trivial_avg_max_prob,
                    "routing_variance": ecc.trivial_routing_variance,
                },
                "cost_optimization": {
                    "safe_cheap_pct": ecc.cost_opt_safe_cheap_pct,
                    "potential_savings": ecc.cost_opt_potential_savings,
                },
            }

        if self.difficulty_pred_corr is not None:
            result["difficulty_pred_corr"] = self.difficulty_pred_corr
        if self.uncertainty_error_corr is not None:
            result["uncertainty_error_corr"] = self.uncertainty_error_corr

        if self.difficulty_detection:
            dd = self.difficulty_detection
            result["difficulty_detection"] = {
                "mean_logit_by_band": dd.mean_logit_by_band,
                "mean_logit_std_by_band": dd.mean_logit_std_by_band,
                "trivial_vs_impossible_auc": dd.trivial_vs_impossible_auc,
                "optimal_threshold": dd.optimal_threshold,
                "separation_accuracy": dd.separation_accuracy,
                "mean_logit_difficulty_corr": dd.mean_logit_difficulty_corr,
                "recommended_thresholds": {
                    "trivial": dd.recommended_trivial_threshold,
                    "hard": dd.recommended_hard_threshold,
                },
                "best_llms_on_hard": dd.best_llms_on_hard,
                "best_llm_success_rates": dd.best_llm_success_rates,
            }

        if self.ranking_discrimination:
            rd = self.ranking_discrimination
            result["ranking_discrimination"] = {
                "avg_spearman_corr": rd.avg_spearman_corr,
                "avg_kendall_tau": rd.avg_kendall_tau,
                "pct_positive_corr": rd.pct_positive_corr,
                "per_llm_discrimination": rd.per_llm_discrimination,
                "avg_llm_discrimination": rd.avg_llm_discrimination,
                "task_affinity_captured": rd.task_affinity_captured,
                "avg_task_affinity": rd.avg_task_affinity,
                "tasks_with_strong_affinity": [
                    str(t) for t in rd.tasks_with_strong_affinity
                ],
                "score_reward_corr": rd.score_reward_corr,
                "high_score_success_rate": rd.high_score_success_rate,
                "low_score_failure_rate": rd.low_score_failure_rate,
            }

        if self.error_analysis:
            ea = self.error_analysis
            result["error_analysis"] = {
                "n_errors": ea.n_errors,
                "error_rate": ea.error_rate,
                "avg_error_regret": ea.avg_error_regret,
                "errors_by_band": ea.errors_by_band,
                "error_rate_by_band": ea.error_rate_by_band,
                "worst_error_regrets": ea.worst_error_regrets,
                "most_confused_pairs": [
                    {"predicted": p, "actual": a, "count": c}
                    for p, a, c in ea.most_confused_pairs
                ],
                "error_gini": ea.error_gini,
                "tasks_with_most_errors": [
                    {"task": str(t), "count": c}
                    for t, c in ea.tasks_with_most_errors
                ],
            }

        if self.cost_quality:
            cq = self.cost_quality
            result["cost_quality"] = {
                "llm_costs": cq.llm_costs,
                "random_routing_cost": cq.random_routing_cost,
                "oracle_routing_cost": cq.oracle_routing_cost,
                "router_routing_cost": cq.router_routing_cost,
                "cheapest_always_cost": cq.cheapest_always_cost,
                "cost_savings_vs_random": cq.cost_savings_vs_random,
                "cost_savings_vs_oracle": cq.cost_savings_vs_oracle,
                "quality_at_budget": cq.quality_at_budget,
                "is_pareto_optimal": cq.is_pareto_optimal,
                "pareto_gap": cq.pareto_gap,
            }

        if self.statistical_significance:
            ss = self.statistical_significance
            result["statistical_significance"] = {
                "frac_of_oracle_ci": list(ss.frac_of_oracle_ci),
                "regret_ci": list(ss.regret_ci),
                "spearman_ci": list(ss.spearman_ci),
                "frac_of_oracle_se": ss.frac_of_oracle_se,
                "regret_se": ss.regret_se,
                "spearman_se": ss.spearman_se,
                "n_samples": ss.n_samples,
                "n_effective": ss.n_effective,
            }

        if self.inference_benchmark:
            ib = self.inference_benchmark
            result["inference_benchmark"] = {
                "latency_mean_ms": ib.latency_mean_ms,
                "latency_p50_ms": ib.latency_p50_ms,
                "latency_p95_ms": ib.latency_p95_ms,
                "latency_p99_ms": ib.latency_p99_ms,
                "queries_per_second": ib.queries_per_second,
                "batch_size_tested": ib.batch_size_tested,
                "model_size_mb": ib.model_size_mb,
                "peak_memory_mb": ib.peak_memory_mb,
                "device": ib.device,
                "dtype": ib.dtype,
            }

        return result


def binary_entropy(p: float) -> float:
    """Compute binary entropy H(p) = -p*log2(p) - (1-p)*log2(1-p)."""
    if p <= 0 or p >= 1:
        return 0.0
    return -p * np.log2(p) - (1 - p) * np.log2(1 - p)


def get_difficulty_band(n_correct: int, n_llms: int) -> str:
    """Map number of correct models to difficulty band."""
    # Normalize to 8-LLM scale if different
    if n_llms != 8:
        n_correct = int(round(n_correct * 8 / n_llms))

    for band_name, (low, high) in DIFFICULTY_BANDS.items():
        if low <= n_correct <= high:
            return band_name
    return "unknown"


class GroundTruthDifficultyAnalyzer:
    """Analyze ground truth data to characterize query difficulty.

    This class computes difficulty metrics from the router_data.csv without
    needing a trained model. It answers questions like:
    - What fraction of queries are easy vs hard?
    - Which tasks are most difficult?
    - Where does model disagreement occur?
    """

    def __init__(
        self,
        router_data_path: str | Path,
        correct_threshold: float = 0.5,
    ) -> None:
        """Initialize analyzer.

        Args:
            router_data_path: Path to router_data.csv or .parquet
            correct_threshold: Threshold for effect to count as correct
        """
        self.router_data_path = Path(router_data_path)
        self.correct_threshold = correct_threshold

        # Load data
        logger.info("Loading router data from %s", self.router_data_path)
        if self.router_data_path.suffix == ".parquet":
            self.df = pd.read_parquet(
                self.router_data_path,
                columns=["task_id", "query", "llm", "effect", "metric"],
            )
        else:
            self.df = pd.read_csv(
                self.router_data_path,
                usecols=["task_id", "query", "llm", "effect", "metric"],
            )

        self.llm_names = self.df["llm"].unique().tolist()
        self.n_llms = len(self.llm_names)
        logger.info("Loaded %d rows, %d LLMs: %s", len(self.df), self.n_llms, self.llm_names)

        # Compute per-query statistics
        self._compute_query_stats()

    def _compute_query_stats(self) -> None:
        """Compute difficulty statistics for each query."""
        df = self.df.copy()
        df["correct"] = (df["effect"] >= self.correct_threshold).astype(int)

        # Group by query
        query_stats = df.groupby("query", sort=False).agg({
            "correct": ["sum", "mean"],
            "effect": ["mean", "std", "min", "max"],
            "task_id": "first",
        }).reset_index()

        query_stats.columns = [
            "query", "n_correct", "pct_correct",
            "effect_mean", "effect_std", "effect_min", "effect_max", "task_id"
        ]

        # Fill NaN std with 0
        query_stats["effect_std"] = query_stats["effect_std"].fillna(0)

        # Compute derived metrics
        query_stats["perf_range"] = query_stats["effect_max"] - query_stats["effect_min"]
        query_stats["agreement_entropy"] = query_stats["pct_correct"].apply(binary_entropy)
        query_stats["difficulty_band"] = query_stats["n_correct"].apply(
            lambda x: get_difficulty_band(int(x), self.n_llms)
        )

        self.query_stats = query_stats
        self.n_queries = len(query_stats)
        logger.info("Computed stats for %d queries", self.n_queries)

    def get_difficulty_distribution(self) -> dict[str, DifficultyBandStats]:
        """Get distribution of queries across difficulty bands."""
        result = {}

        for band_name in DIFFICULTY_BANDS.keys():
            mask = self.query_stats["difficulty_band"] == band_name
            band_df = self.query_stats[mask]

            if len(band_df) == 0:
                continue

            result[band_name] = DifficultyBandStats(
                band_name=band_name,
                n_queries=len(band_df),
                pct_of_total=len(band_df) / self.n_queries * 100,
                avg_effect=band_df["effect_mean"].mean(),
                avg_effect_std=band_df["effect_std"].mean(),
                avg_perf_range=band_df["perf_range"].mean(),
                avg_entropy=band_df["agreement_entropy"].mean(),
            )

        return result

    def get_task_difficulty(self) -> dict[str, float]:
        """Get average difficulty (n_correct) per task."""
        task_stats = self.query_stats.groupby("task_id")["n_correct"].mean()
        return task_stats.sort_values().to_dict()

    def get_high_routing_value_queries(self, top_k: int = 100) -> pd.DataFrame:
        """Get queries with highest routing value (perf_range)."""
        return self.query_stats.nlargest(top_k, "perf_range")[
            ["query", "task_id", "n_correct", "perf_range", "effect_min", "effect_max"]
        ]

    def get_reward_matrix(self) -> tuple[np.ndarray, list[str], list[str]]:
        """Get query x LLM reward matrix for routing evaluation.

        Returns:
            rewards: (n_queries, n_llms) array of effects
            query_ids: List of query strings
            llm_names: List of LLM names
        """
        # Pivot to get query x llm matrix
        pivot = self.df.pivot_table(
            index="query",
            columns="llm",
            values="effect",
            aggfunc="first",
        )

        # Ensure consistent LLM ordering
        llm_names = [llm for llm in self.llm_names if llm in pivot.columns]
        pivot = pivot[llm_names]

        # Drop rows with NaN (incomplete coverage)
        pivot = pivot.dropna()

        return pivot.values, pivot.index.tolist(), llm_names

    def print_summary(self) -> None:
        """Print summary of ground truth difficulty analysis."""
        print("\n" + "=" * 70)
        print("GROUND TRUTH DIFFICULTY ANALYSIS")
        print("=" * 70)
        print(f"Total queries: {self.n_queries:,}")
        print(f"LLMs: {self.n_llms} ({', '.join(self.llm_names)})")

        # Difficulty distribution
        print("\n--- Difficulty Distribution ---")
        dist = self.get_difficulty_distribution()
        for band_name in ["trivial", "easy", "medium", "hard", "very_hard", "impossible"]:
            if band_name not in dist:
                continue
            stats = dist[band_name]
            bar = "█" * int(stats.pct_of_total / 2)
            print(
                f"  {band_name:12s}: {stats.n_queries:5,} ({stats.pct_of_total:5.1f}%) "
                f"perf_range={stats.avg_perf_range:.2f} {bar}"
            )

        # Task difficulty
        print("\n--- Task Difficulty (avg models correct) ---")
        task_diff = self.get_task_difficulty()
        for task_id, avg_correct in list(task_diff.items())[:10]:
            print(f"  {task_id:30s}: {avg_correct:.2f}")
        if len(task_diff) > 10:
            print(f"  ... and {len(task_diff) - 10} more tasks")

        # Routing value
        print("\n--- Highest Routing Value Queries ---")
        high_rv = self.get_high_routing_value_queries(5)
        for _, row in high_rv.iterrows():
            print(
                f"  [{row['task_id']}] perf_range={row['perf_range']:.2f} "
                f"(min={row['effect_min']:.2f}, max={row['effect_max']:.2f})"
            )
            print(f"    {row['query'][:60]}...")


def compute_routing_metrics(
    predictions: np.ndarray,
    rewards: np.ndarray,
) -> RoutingMetrics:
    """Compute routing metrics from predictions and rewards.

    Args:
        predictions: (N, num_llms) - predicted scores/probabilities
        rewards: (N, num_llms) - ground truth rewards/effects

    Returns:
        RoutingMetrics with all computed values
    """
    predictions = torch.tensor(predictions, dtype=torch.float32)
    rewards = torch.tensor(rewards, dtype=torch.float32)

    pred_rank = torch.argsort(predictions, dim=1, descending=True)
    true_rank = torch.argsort(rewards, dim=1, descending=True)

    max_reward = rewards.max(dim=1).values
    min_reward = rewards.min(dim=1).values
    span = max_reward - min_reward
    shifted_rewards = rewards - min_reward.unsqueeze(1)

    row_indices = torch.arange(rewards.size(0))
    chosen_reward = rewards[row_indices, pred_rank[:, 0]]
    chosen_shifted = shifted_rewards[row_indices, pred_rank[:, 0]]
    oracle_shifted = shifted_rewards[row_indices, true_rank[:, 0]]

    top1_acc = (pred_rank[:, 0] == true_rank[:, 0]).float().mean().item()
    regret = (max_reward - chosen_reward).mean().item()

    # frac_of_oracle with correct tie handling.
    # For ties (all-correct or all-wrong), any choice = oracle -> frac = 1.0.
    is_tie = span < 1e-6
    per_query_frac = torch.where(
        is_tie,
        torch.ones_like(chosen_shifted),
        chosen_shifted / oracle_shifted.clamp_min(1e-6),
    )
    frac_of_oracle = per_query_frac.mean().item()

    # NDCG@k
    def ndcg_at_k(k: int) -> float:
        k = min(k, rewards.size(1))
        denom = torch.log2(torch.arange(2, k + 2, dtype=torch.float32))
        gains_pred = shifted_rewards[row_indices.unsqueeze(1), pred_rank[:, :k]]
        gains_true = shifted_rewards[row_indices.unsqueeze(1), true_rank[:, :k]]
        dcg = (gains_pred / denom).sum(dim=1)
        idcg = (gains_true / denom).sum(dim=1) + 1e-8
        return (dcg / idcg).mean().item()

    return RoutingMetrics(
        top1_acc=top1_acc,
        regret=regret,
        frac_of_oracle=frac_of_oracle,
        ndcg_at_1=ndcg_at_k(1),
        ndcg_at_3=ndcg_at_k(3),
        avg_chosen_reward=chosen_reward.mean().item(),
        avg_oracle_reward=max_reward.mean().item(),
    )


def compute_calibration_metrics(
    predictions: np.ndarray,
    rewards: np.ndarray,
    n_bins: int = 10,
    correct_threshold: float = 0.5,
) -> CalibrationMetrics:
    """Compute uncertainty calibration metrics.

    This function tests the "knows when it doesn't know" hypothesis by measuring:
    1. Whether entropy correlates with difficulty (entropy_difficulty_corr)
    2. Whether entropy predicts routing errors (entropy_error_corr)
    3. Whether entropy predicts routing quality (entropy_quality_corr)

    A well-calibrated uncertainty signal should have:
    - Positive entropy_error_corr (high entropy -> model makes more mistakes)
    - Negative entropy_quality_corr (high entropy -> lower routing quality)

    Args:
        predictions: (N, num_llms) - predicted logits
        rewards: (N, num_llms) - ground truth effects
        n_bins: Number of bins for ECE computation
        correct_threshold: Threshold for correct classification

    Returns:
        CalibrationMetrics with ECE, MCE, entropy correlations, and distribution
    """
    n_queries, n_llms = predictions.shape

    # Softmax predictions to get probabilities
    probs = torch.softmax(torch.tensor(predictions), dim=1).numpy()

    # Max probability (confidence) and prediction
    max_prob = probs.max(axis=1)
    pred_idx = probs.argmax(axis=1)

    # Ground truth: which LLM has highest reward
    true_idx = rewards.argmax(axis=1)
    correct = (pred_idx == true_idx).astype(float)
    errors = 1 - correct  # 1 if routing error, 0 if correct

    # ECE: Expected Calibration Error
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    mce = 0.0

    for i in range(n_bins):
        in_bin = (max_prob >= bin_boundaries[i]) & (max_prob < bin_boundaries[i + 1])
        n_in_bin = in_bin.sum()

        if n_in_bin > 0:
            avg_confidence = max_prob[in_bin].mean()
            avg_accuracy = correct[in_bin].mean()
            gap = abs(avg_accuracy - avg_confidence)
            ece += (n_in_bin / len(max_prob)) * gap
            mce = max(mce, gap)

    # Overconfidence: P(wrong | high_confidence)
    high_conf_mask = max_prob >= 0.8
    if high_conf_mask.sum() > 0:
        overconfidence_rate = 1 - correct[high_conf_mask].mean()
    else:
        overconfidence_rate = 0.0

    # Underconfidence: P(correct | low_confidence)
    low_conf_mask = max_prob < 0.5
    if low_conf_mask.sum() > 0:
        underconfidence_rate = correct[low_conf_mask].mean()
    else:
        underconfidence_rate = 0.0

    # Prediction entropy (normalized to [0, 1])
    entropy = -np.sum(probs * np.log(probs + 1e-8), axis=1)
    max_entropy = np.log(n_llms)
    norm_entropy = entropy / max_entropy

    # Ground truth difficulty: number of models correct
    n_correct = (rewards >= correct_threshold).sum(axis=1)

    # Per-query routing quality (frac_of_oracle)
    chosen_reward = rewards[np.arange(n_queries), pred_idx]
    oracle_reward = rewards.max(axis=1)
    min_reward = rewards.min(axis=1)
    # Handle division by zero for queries where all LLMs have same performance
    denominator = oracle_reward - min_reward
    valid_denom = denominator > 1e-6
    frac_of_oracle = np.ones(n_queries)  # Default to 1.0 (perfect) for degenerate cases
    frac_of_oracle[valid_denom] = (
        (chosen_reward[valid_denom] - min_reward[valid_denom]) /
        denominator[valid_denom]
    )

    # === CORRELATIONS: Does entropy capture "knows when it doesn't know"? ===

    # 1. Entropy-difficulty correlation: Does entropy reflect query difficulty?
    entropy_difficulty_corr = np.corrcoef(norm_entropy, n_correct)[0, 1]

    # 2. Confidence-difficulty correlation
    confidence_difficulty_corr = np.corrcoef(max_prob, n_correct)[0, 1]

    # 3. Entropy-error correlation: Does high entropy predict routing mistakes?
    # A positive correlation means the model "knows when it doesn't know"
    entropy_error_corr = np.corrcoef(norm_entropy, errors)[0, 1]

    # 4. Entropy-quality correlation: Does low entropy predict high routing quality?
    # A negative correlation means the model is confident when it's good
    entropy_quality_corr = np.corrcoef(norm_entropy, frac_of_oracle)[0, 1]

    # === ENTROPY DISTRIBUTION: Is entropy actually informative? ===
    avg_entropy = float(norm_entropy.mean())
    entropy_std = float(norm_entropy.std())
    pct_high_entropy = float((norm_entropy > 0.9).mean() * 100)

    # === RAW LOGIT UNCERTAINTY: Bypass softmax compression ===
    # Softmax with small logit differences produces near-uniform probabilities,
    # crushing the uncertainty signal. Instead, compute uncertainty from raw logits:

    # logit_std: High std = model has strong preferences (confident)
    logit_std = predictions.std(axis=1)

    # logit_margin: max - second_max (how much better is top choice?)
    sorted_logits = np.sort(predictions, axis=1)[:, ::-1]  # descending
    logit_margin = sorted_logits[:, 0] - sorted_logits[:, 1]

    # Correlations: High std/margin should predict FEWER errors, HIGHER quality
    logit_std_error_corr = np.corrcoef(logit_std, errors)[0, 1]
    logit_std_quality_corr = np.corrcoef(logit_std, frac_of_oracle)[0, 1]
    logit_margin_error_corr = np.corrcoef(logit_margin, errors)[0, 1]
    logit_margin_quality_corr = np.corrcoef(logit_margin, frac_of_oracle)[0, 1]

    avg_logit_std = float(logit_std.mean())
    avg_logit_margin = float(logit_margin.mean())

    return CalibrationMetrics(
        ece=ece,
        mce=mce,
        overconfidence_rate=overconfidence_rate,
        underconfidence_rate=underconfidence_rate,
        entropy_difficulty_corr=entropy_difficulty_corr,
        confidence_difficulty_corr=confidence_difficulty_corr,
        entropy_error_corr=entropy_error_corr,
        entropy_quality_corr=entropy_quality_corr,
        avg_entropy=avg_entropy,
        entropy_std=entropy_std,
        pct_high_entropy=pct_high_entropy,
        logit_std_error_corr=logit_std_error_corr,
        logit_std_quality_corr=logit_std_quality_corr,
        logit_margin_error_corr=logit_margin_error_corr,
        logit_margin_quality_corr=logit_margin_quality_corr,
        avg_logit_std=avg_logit_std,
        avg_logit_margin=avg_logit_margin,
    )


def compute_routing_by_difficulty(
    predictions: np.ndarray,
    rewards: np.ndarray,
    query_difficulty: np.ndarray,  # n_correct per query
    n_llms: int = 8,
) -> dict[str, RoutingMetrics]:
    """Compute routing metrics stratified by difficulty band.

    Args:
        predictions: (N, num_llms) predictions
        rewards: (N, num_llms) rewards
        query_difficulty: (N,) number of models correct
        n_llms: Number of LLMs for difficulty band calculation

    Returns:
        Dict mapping band name to RoutingMetrics
    """
    result = {}

    for band_name, (low, high) in DIFFICULTY_BANDS.items():
        # Get mask for this band
        if n_llms != 8:
            # Normalize to 8-LLM scale
            normalized = query_difficulty * 8 / n_llms
            mask = (normalized >= low) & (normalized <= high)
        else:
            mask = (query_difficulty >= low) & (query_difficulty <= high)

        if mask.sum() < 10:  # Skip if too few samples
            continue

        band_preds = predictions[mask]
        band_rewards = rewards[mask]

        result[band_name] = compute_routing_metrics(band_preds, band_rewards)

    return result


def compute_edge_case_calibration(
    predictions: np.ndarray,
    rewards: np.ndarray,
    n_llms: int,
    correct_threshold: float = 0.5,
    cheap_model_idx: int = 0,  # Index of cheapest model
    cheap_cost_ratio: float = 0.015,  # Cost ratio cheap/expensive (e.g., Nano vs Opus)
) -> EdgeCaseCalibration:
    """Compute calibration metrics for impossible and trivial queries.

    These edge cases are critical for understanding whether the model
    "knows what it doesn't know":
    - Impossible queries: Should show high uncertainty (no good choice exists)
    - Trivial queries: Should show low uncertainty (all choices are good)

    Args:
        predictions: (N, num_llms) - predicted logits
        rewards: (N, num_llms) - ground truth effects
        n_llms: Number of LLMs
        correct_threshold: Threshold for counting as correct
        cheap_model_idx: Index of the cheapest model for cost optimization
        cheap_cost_ratio: Cost ratio of cheap model vs average

    Returns:
        EdgeCaseCalibration with metrics for impossible/trivial queries
    """
    # Convert to probabilities
    probs = torch.softmax(torch.tensor(predictions), dim=1).numpy()
    max_prob = probs.max(axis=1)

    # Compute entropy (normalized to [0, 1])
    entropy = -np.sum(probs * np.log(probs + 1e-8), axis=1)
    max_entropy = np.log(n_llms)
    norm_entropy = entropy / max_entropy

    # Ground truth difficulty
    n_correct = (rewards >= correct_threshold).sum(axis=1)

    # Identify impossible and trivial queries
    impossible_mask = n_correct == 0
    trivial_mask = n_correct == n_llms

    n_impossible = impossible_mask.sum()
    n_trivial = trivial_mask.sum()
    n_total = len(predictions)

    # Impossible query metrics
    if n_impossible > 0:
        impossible_avg_entropy = norm_entropy[impossible_mask].mean()
        impossible_avg_max_prob = max_prob[impossible_mask].mean()
        # Overconfidence: predicting with high confidence when all models fail
        impossible_overconfidence = (max_prob[impossible_mask] > 0.7).mean()
    else:
        impossible_avg_entropy = 0.0
        impossible_avg_max_prob = 0.0
        impossible_overconfidence = 0.0

    # Trivial query metrics
    if n_trivial > 0:
        trivial_avg_entropy = norm_entropy[trivial_mask].mean()
        trivial_avg_max_prob = max_prob[trivial_mask].mean()
        # Routing variance: how much do predictions vary when it doesn't matter?
        trivial_routing_variance = probs[trivial_mask].var(axis=1).mean()
    else:
        trivial_avg_entropy = 0.0
        trivial_avg_max_prob = 0.0
        trivial_routing_variance = 0.0

    # Cost optimization: Can we safely route edge cases to cheap models?
    # "Safe" means the cheap model performs as well as the predicted best
    edge_case_mask = impossible_mask | trivial_mask
    n_edge_cases = edge_case_mask.sum()

    if n_edge_cases > 0:
        # For impossible: all models fail, so cheap is fine
        # For trivial: all models succeed, so cheap is fine
        # "Safe" = cheap model's reward >= 0.9 * max reward on that query
        cheap_rewards = rewards[edge_case_mask, cheap_model_idx]
        max_rewards = rewards[edge_case_mask].max(axis=1)

        safe_mask = cheap_rewards >= 0.9 * max_rewards
        cost_opt_safe_cheap_pct = safe_mask.mean()

        # Potential savings: edge_case_pct * (1 - cheap_cost_ratio)
        edge_case_pct = n_edge_cases / n_total
        cost_opt_potential_savings = edge_case_pct * (1 - cheap_cost_ratio)
    else:
        cost_opt_safe_cheap_pct = 0.0
        cost_opt_potential_savings = 0.0

    return EdgeCaseCalibration(
        n_impossible=int(n_impossible),
        pct_impossible=100 * n_impossible / n_total,
        impossible_avg_entropy=float(impossible_avg_entropy),
        impossible_avg_max_prob=float(impossible_avg_max_prob),
        impossible_overconfidence=float(impossible_overconfidence),
        n_trivial=int(n_trivial),
        pct_trivial=100 * n_trivial / n_total,
        trivial_avg_entropy=float(trivial_avg_entropy),
        trivial_avg_max_prob=float(trivial_avg_max_prob),
        trivial_routing_variance=float(trivial_routing_variance),
        cost_opt_safe_cheap_pct=float(cost_opt_safe_cheap_pct),
        cost_opt_potential_savings=float(cost_opt_potential_savings),
    )


def compute_difficulty_detection_metrics(
    mean_logit: np.ndarray,
    rewards: np.ndarray,
    llm_names: list[str],
    correct_threshold: float = 0.5,
) -> DifficultyDetection:
    """Compute difficulty detection metrics using mean logit signal.

    Key insight: The mean logit across LLMs captures query difficulty:
    - High mean logit: Model believes most LLMs will succeed (trivial query)
    - Low/negative mean logit: Model believes most LLMs will fail (hard query)

    This provides a production-ready signal for:
    1. Routing trivial queries to cheap models
    2. Detecting impossible queries for cost optimization
    3. Identifying hard queries for fallback to best-performing LLMs

    Args:
        mean_logit: (N,) - mean logit per query from model predictions
        rewards: (N, num_llms) - ground truth rewards/effects
        llm_names: List of LLM names in order
        correct_threshold: Threshold for counting as correct

    Returns:
        DifficultyDetection with all computed metrics
    """
    n_queries, n_llms = rewards.shape

    # Ground truth difficulty
    n_correct = (rewards >= correct_threshold).sum(axis=1)

    # Assign difficulty bands
    def get_band(nc: int) -> str:
        if nc == 0:
            return "impossible"
        elif nc <= 2:
            return "very_hard"
        elif nc <= 4:
            return "hard"
        elif nc <= 6:
            return "medium"
        else:
            return "trivial"

    bands = np.array([get_band(nc) for nc in n_correct])

    # Mean logit statistics by band
    mean_logit_by_band = {}
    mean_logit_std_by_band = {}
    for band in ["trivial", "medium", "hard", "very_hard", "impossible"]:
        mask = bands == band
        if mask.sum() > 0:
            mean_logit_by_band[band] = float(mean_logit[mask].mean())
            mean_logit_std_by_band[band] = float(mean_logit[mask].std())

    # Correlation with ground truth difficulty
    mean_logit_difficulty_corr = float(np.corrcoef(mean_logit, n_correct)[0, 1])

    # Trivial vs impossible separation
    trivial_mask = bands == "trivial"
    impossible_mask = bands == "impossible"

    if trivial_mask.sum() > 0 and impossible_mask.sum() > 0:
        trivial_ml = mean_logit[trivial_mask]
        impossible_ml = mean_logit[impossible_mask]

        # Find optimal threshold
        all_ml = np.concatenate([trivial_ml, impossible_ml])
        all_labels = np.concatenate([
            np.ones(len(trivial_ml)),
            np.zeros(len(impossible_ml)),
        ])

        best_acc = 0.0
        best_thresh = 0.0
        for thresh in np.linspace(all_ml.min(), all_ml.max(), 100):
            pred = (all_ml > thresh).astype(int)
            acc = (pred == all_labels).mean()
            if acc > best_acc:
                best_acc = acc
                best_thresh = thresh

        # Compute AUC
        correct_pairs = sum(
            1 for t in trivial_ml for i in impossible_ml if t > i
        )
        total_pairs = len(trivial_ml) * len(impossible_ml)
        auc = correct_pairs / total_pairs if total_pairs > 0 else 0.5

        trivial_vs_impossible_auc = float(auc)
        optimal_threshold = float(best_thresh)
        separation_accuracy = float(best_acc)
    else:
        trivial_vs_impossible_auc = 0.5
        optimal_threshold = 0.0
        separation_accuracy = 0.5

    # Recommended thresholds (based on band means with safety margin)
    trivial_mean = mean_logit_by_band.get("trivial", 0.5)
    medium_mean = mean_logit_by_band.get("medium", 0.0)
    very_hard_mean = mean_logit_by_band.get("very_hard", -0.5)
    impossible_mean = mean_logit_by_band.get("impossible", -0.3)

    # Threshold for trivial: midpoint between trivial and medium
    recommended_trivial_threshold = (trivial_mean + medium_mean) / 2
    # Threshold for hard: midpoint between hard and very_hard
    hard_mean = mean_logit_by_band.get("hard", 0.0)
    recommended_hard_threshold = (hard_mean + very_hard_mean) / 2

    # Best LLMs on very_hard queries
    # Note: rewards may have more LLMs than the model - only use model's LLMs
    num_model_llms = len(llm_names)
    very_hard_mask = bands == "very_hard"
    if very_hard_mask.sum() > 0:
        vh_rewards = rewards[very_hard_mask][:, :num_model_llms]
        vh_success = (vh_rewards >= correct_threshold).mean(axis=0)
        sorted_idx = np.argsort(vh_success)[::-1]
        best_llms_on_hard = [llm_names[i] for i in sorted_idx[:3]]
        best_llm_success_rates = {
            llm_names[i]: float(vh_success[i]) for i in sorted_idx
        }
    else:
        best_llms_on_hard = llm_names[:3]
        best_llm_success_rates = {name: 0.0 for name in llm_names}

    return DifficultyDetection(
        mean_logit_by_band=mean_logit_by_band,
        mean_logit_std_by_band=mean_logit_std_by_band,
        trivial_vs_impossible_auc=trivial_vs_impossible_auc,
        optimal_threshold=optimal_threshold,
        separation_accuracy=separation_accuracy,
        mean_logit_difficulty_corr=mean_logit_difficulty_corr,
        recommended_trivial_threshold=float(recommended_trivial_threshold),
        recommended_hard_threshold=float(recommended_hard_threshold),
        best_llms_on_hard=best_llms_on_hard,
        best_llm_success_rates=best_llm_success_rates,
    )


def compute_ranking_discrimination(
    predictions: np.ndarray,
    rewards: np.ndarray,
    llm_names: list[str],
    task_ids: np.ndarray | None = None,
    correct_threshold: float = 0.5,
) -> RankingDiscrimination:
    """Compute ranking discrimination metrics.

    These metrics answer: "Are the router's rankings meaningful?"

    Key questions tested:
    1. Per-query ranking correlation: Do predicted ranks match true ranks?
    2. Per-LLM discrimination: When router gives high score, is LLM actually better?
    3. Task affinity: Does router learn which LLMs excel on which tasks?
    4. Score-reward calibration: Overall correlation between scores and rewards.

    Args:
        predictions: (N, num_llms) - predicted logits/scores
        rewards: (N, num_llms) - ground truth effects/rewards
        llm_names: List of LLM names in order
        task_ids: (N,) - task ID per query (optional, for task affinity analysis)
        correct_threshold: Threshold for "correct" classification

    Returns:
        RankingDiscrimination with all discrimination metrics
    """
    from scipy.stats import kendalltau, spearmanr

    n_queries, n_pred_llms = predictions.shape
    n_reward_llms = rewards.shape[1]

    # Handle mismatch: model may have fewer LLMs than dataset
    n_llms = min(n_pred_llms, n_reward_llms)
    if n_pred_llms != n_reward_llms:
        logger.warning(
            "LLM count mismatch: predictions=%d, rewards=%d. Using first %d.",
            n_pred_llms, n_reward_llms, n_llms
        )
        predictions = predictions[:, :n_llms]
        rewards = rewards[:, :n_llms]
        llm_names = llm_names[:n_llms]

    # === PER-QUERY RANKING CORRELATION ===
    # For each query, compute correlation between predicted and true ranks
    spearman_corrs = []
    kendall_taus = []

    for i in range(n_queries):
        pred_row = predictions[i]
        true_row = rewards[i]

        # Skip if all rewards are the same (no ranking possible)
        if true_row.std() < 1e-6:
            continue

        sp_corr, _ = spearmanr(pred_row, true_row)
        kt_corr, _ = kendalltau(pred_row, true_row)

        if not np.isnan(sp_corr):
            spearman_corrs.append(sp_corr)
        if not np.isnan(kt_corr):
            kendall_taus.append(kt_corr)

    avg_spearman = float(np.mean(spearman_corrs)) if spearman_corrs else 0.0
    avg_kendall = float(np.mean(kendall_taus)) if kendall_taus else 0.0
    pct_positive = float(np.mean([c > 0 for c in spearman_corrs]) * 100) if spearman_corrs else 0.0

    # === PER-LLM DISCRIMINATION ===
    # For each LLM: correlation between router's score and actual reward
    per_llm_disc = {}
    for j, llm_name in enumerate(llm_names):
        llm_scores = predictions[:, j]
        llm_rewards = rewards[:, j]

        # Correlation: when router gives high score, is reward actually high?
        corr = np.corrcoef(llm_scores, llm_rewards)[0, 1]
        per_llm_disc[llm_name] = float(corr) if not np.isnan(corr) else 0.0

    avg_llm_disc = float(np.mean(list(per_llm_disc.values())))

    # === TASK-SPECIFIC AFFINITY ===
    # For each task: does router capture which LLMs are best?
    task_affinity = {}
    if task_ids is not None:
        unique_tasks = np.unique(task_ids)
        for task in unique_tasks:
            task_mask = task_ids == task
            if task_mask.sum() < 10:
                continue

            task_preds = predictions[task_mask]
            task_rewards = rewards[task_mask]

            # Average predictions and rewards per LLM for this task
            avg_pred_per_llm = task_preds.mean(axis=0)
            avg_reward_per_llm = task_rewards.mean(axis=0)

            # Does router's avg preference match actual avg performance?
            corr = np.corrcoef(avg_pred_per_llm, avg_reward_per_llm)[0, 1]
            task_affinity[task] = float(corr) if not np.isnan(corr) else 0.0

    avg_task_aff = float(np.mean(list(task_affinity.values()))) if task_affinity else 0.0
    strong_affinity_tasks = [t for t, c in task_affinity.items() if c > 0.5]

    # === OVERALL SCORE-REWARD CALIBRATION ===
    # Flatten all scores and rewards
    all_scores = predictions.flatten()
    all_rewards = rewards.flatten()
    score_reward_corr = float(np.corrcoef(all_scores, all_rewards)[0, 1])

    # High score success rate: P(correct | score in top 25%)
    score_threshold_high = np.percentile(all_scores, 75)
    score_threshold_low = np.percentile(all_scores, 25)

    high_score_mask = all_scores >= score_threshold_high
    low_score_mask = all_scores <= score_threshold_low

    correct = (all_rewards >= correct_threshold)
    high_score_success = float(correct[high_score_mask].mean()) if high_score_mask.sum() > 0 else 0.0
    low_score_failure = float((~correct[low_score_mask]).mean()) if low_score_mask.sum() > 0 else 0.0

    return RankingDiscrimination(
        avg_spearman_corr=avg_spearman,
        avg_kendall_tau=avg_kendall,
        pct_positive_corr=pct_positive,
        per_llm_discrimination=per_llm_disc,
        avg_llm_discrimination=avg_llm_disc,
        task_affinity_captured=task_affinity,
        avg_task_affinity=avg_task_aff,
        tasks_with_strong_affinity=strong_affinity_tasks,
        score_reward_corr=score_reward_corr,
        high_score_success_rate=high_score_success,
        low_score_failure_rate=low_score_failure,
    )


def compute_error_analysis(
    predictions: np.ndarray,
    rewards: np.ndarray,
    llm_names: list[str],
    task_ids: np.ndarray | None = None,
    correct_threshold: float = 0.5,
) -> ErrorAnalysis:
    """Compute detailed error analysis.

    Answers: "Where and why does the router fail?"

    Args:
        predictions: (N, num_llms) - predicted logits
        rewards: (N, num_llms) - ground truth rewards
        llm_names: List of LLM names
        task_ids: (N,) - task ID per query (optional)
        correct_threshold: Threshold for "correct" classification

    Returns:
        ErrorAnalysis with detailed failure mode information
    """
    n_queries, n_llms = predictions.shape

    # Handle LLM count mismatch
    n_llms = min(n_llms, rewards.shape[1], len(llm_names))
    predictions = predictions[:, :n_llms]
    rewards = rewards[:, :n_llms]
    llm_names = llm_names[:n_llms]

    # Predictions and ground truth
    pred_idx = predictions.argmax(axis=1)
    true_idx = rewards.argmax(axis=1)

    # Error identification
    is_error = pred_idx != true_idx
    n_errors = int(is_error.sum())
    error_rate = n_errors / n_queries

    # Regret calculation
    chosen_reward = rewards[np.arange(n_queries), pred_idx]
    oracle_reward = rewards.max(axis=1)
    regret = oracle_reward - chosen_reward

    avg_error_regret = float(regret[is_error].mean()) if n_errors > 0 else 0.0

    # Errors by difficulty band
    n_correct = (rewards >= correct_threshold).sum(axis=1)
    errors_by_band = {}
    error_rate_by_band = {}

    for band_name, (low, high) in DIFFICULTY_BANDS.items():
        band_mask = (n_correct >= low) & (n_correct <= high)
        n_in_band = band_mask.sum()
        if n_in_band > 0:
            errors_in_band = is_error[band_mask].sum()
            errors_by_band[band_name] = int(errors_in_band)
            error_rate_by_band[band_name] = float(errors_in_band / n_in_band)

    # Worst errors (highest regret)
    error_indices = np.where(is_error)[0]
    if len(error_indices) > 0:
        error_regrets = regret[error_indices]
        sorted_idx = np.argsort(error_regrets)[::-1][:10]
        worst_error_regrets = [float(error_regrets[i]) for i in sorted_idx]
        worst_error_indices = [int(error_indices[i]) for i in sorted_idx]
    else:
        worst_error_regrets = []
        worst_error_indices = []

    # Confusion matrix
    confusion_matrix: dict[str, dict[str, int]] = {name: {} for name in llm_names}
    for i in range(n_queries):
        if is_error[i]:
            pred_name = llm_names[pred_idx[i]]
            true_name = llm_names[true_idx[i]]
            if true_name not in confusion_matrix[pred_name]:
                confusion_matrix[pred_name][true_name] = 0
            confusion_matrix[pred_name][true_name] += 1

    # Most confused pairs
    confused_pairs = []
    for pred_name, actuals in confusion_matrix.items():
        for true_name, count in actuals.items():
            confused_pairs.append((pred_name, true_name, count))
    confused_pairs.sort(key=lambda x: -x[2])
    most_confused_pairs = confused_pairs[:5]

    # Error concentration by task
    tasks_with_most_errors: list[tuple[str, int]] = []
    error_gini = 0.0

    if task_ids is not None:
        task_errors: dict[str, int] = {}
        for i, task in enumerate(task_ids):
            task = str(task)
            if task not in task_errors:
                task_errors[task] = 0
            if is_error[i]:
                task_errors[task] += 1

        sorted_tasks = sorted(task_errors.items(), key=lambda x: -x[1])
        tasks_with_most_errors = sorted_tasks[:5]

        # Gini coefficient
        error_counts = np.array(list(task_errors.values()))
        if len(error_counts) > 1 and error_counts.sum() > 0:
            sorted_counts = np.sort(error_counts)
            n = len(sorted_counts)
            cumsum = np.cumsum(sorted_counts)
            error_gini = float((2 * np.sum((np.arange(1, n + 1) * sorted_counts)) - (n + 1) * cumsum[-1]) / (n * cumsum[-1]))

    return ErrorAnalysis(
        n_errors=n_errors,
        error_rate=error_rate,
        avg_error_regret=avg_error_regret,
        errors_by_band=errors_by_band,
        error_rate_by_band=error_rate_by_band,
        worst_error_regrets=worst_error_regrets,
        worst_error_indices=worst_error_indices,
        confusion_matrix=confusion_matrix,
        most_confused_pairs=most_confused_pairs,
        error_gini=error_gini,
        tasks_with_most_errors=tasks_with_most_errors,
    )


def compute_cost_quality_analysis(
    predictions: np.ndarray,
    rewards: np.ndarray,
    llm_names: list[str],
    llm_costs: dict[str, float] | None = None,
) -> CostQualityAnalysis:
    """Compute cost-quality tradeoff analysis.

    Answers: "What's the dollar impact of this router?"

    Args:
        predictions: (N, num_llms) - predicted logits
        rewards: (N, num_llms) - ground truth rewards
        llm_names: List of LLM names
        llm_costs: Dict of LLM -> relative cost (if None, uses defaults)

    Returns:
        CostQualityAnalysis with cost-quality tradeoffs
    """
    n_queries, n_llms = predictions.shape

    # Default cost model (relative to cheapest)
    if llm_costs is None:
        llm_costs = {
            "GPT-5 Nano": 1.0,
            "Gemini 2.5 Flash Lite": 1.5,
            "Claude Haiku 4.5": 2.0,
            "GPT-5.1": 10.0,
            "GPT-5.1 Codex": 15.0,
            "Gemini 3 Flash Preview": 5.0,
            "Gemini 3 Pro Preview": 20.0,
            "Claude Opus 4.5": 50.0,
        }

    # Build cost vector aligned with llm_names
    costs = np.array([llm_costs.get(name, 5.0) for name in llm_names[:n_llms]])

    # Predictions and ground truth
    pred_idx = predictions.argmax(axis=1)
    true_idx = rewards.argmax(axis=1)
    cheapest_idx = costs.argmin()

    # Cost calculations
    router_cost = costs[pred_idx].mean()
    oracle_cost = costs[true_idx].mean()
    random_cost = costs.mean()
    cheapest_cost = float(costs[cheapest_idx])

    # Cost savings
    cost_savings_vs_random = (random_cost - router_cost) / random_cost if random_cost > 0 else 0
    cost_savings_vs_oracle = 1 - (router_cost / oracle_cost) if oracle_cost > 0 else 0

    # Quality at different cost budgets
    quality_at_budget = {}
    for budget_name, budget_multiplier in [("cheap", 0.5), ("medium", 1.0), ("expensive", 2.0)]:
        budget = random_cost * budget_multiplier
        # Filter to LLMs within budget
        within_budget = costs <= budget
        if within_budget.any():
            # Use router predictions but only among affordable LLMs
            budget_preds = predictions.copy()
            budget_preds[:, ~within_budget] = -np.inf
            budget_pred_idx = budget_preds.argmax(axis=1)

            chosen = rewards[np.arange(n_queries), budget_pred_idx]
            oracle = rewards.max(axis=1)
            minimum = rewards.min(axis=1)
            span = oracle - minimum
            # Handle ties: for all-correct or all-wrong, any choice = oracle -> frac = 1.0
            is_tie = span < 1e-6
            per_query_frac = np.where(
                is_tie,
                1.0,
                (chosen - minimum) / np.maximum(span, 1e-6),
            )
            frac = per_query_frac.mean()
            quality_at_budget[budget_name] = float(frac)

    # Pareto optimality (simplified check)
    # A router is Pareto optimal if no other achieves same quality at lower cost
    # Here we just check if we're close to the quality-cost frontier
    is_pareto_optimal = bool(
        cost_savings_vs_random > 0 or router_cost <= oracle_cost * 1.1
    )
    pareto_gap = max(0, (router_cost - oracle_cost) / random_cost) if random_cost > 0 else 0

    return CostQualityAnalysis(
        llm_costs=llm_costs,
        random_routing_cost=float(random_cost),
        oracle_routing_cost=float(oracle_cost),
        router_routing_cost=float(router_cost),
        cheapest_always_cost=cheapest_cost,
        cost_savings_vs_random=float(cost_savings_vs_random),
        cost_savings_vs_oracle=float(cost_savings_vs_oracle),
        quality_at_budget=quality_at_budget,
        is_pareto_optimal=is_pareto_optimal,
        pareto_gap=float(pareto_gap),
    )


def compute_statistical_significance(
    predictions: np.ndarray,
    rewards: np.ndarray,
    n_bootstrap: int = 1000,
) -> StatisticalSignificance:
    """Compute confidence intervals via bootstrap.

    Answers: "Is this difference real or just noise?"

    Args:
        predictions: (N, num_llms) - predicted logits
        rewards: (N, num_llms) - ground truth rewards
        n_bootstrap: Number of bootstrap samples

    Returns:
        StatisticalSignificance with confidence intervals
    """
    from scipy.stats import spearmanr

    n_queries = predictions.shape[0]

    def compute_metrics(indices: np.ndarray) -> tuple[float, float, float]:
        """Compute frac_of_oracle, regret, spearman for given indices."""
        preds = predictions[indices]
        rews = rewards[indices]

        pred_idx = preds.argmax(axis=1)
        true_idx = rews.argmax(axis=1)

        chosen = rews[np.arange(len(indices)), pred_idx]
        oracle = rews.max(axis=1)
        minimum = rews.min(axis=1)

        regret = (oracle - chosen).mean()

        # Handle ties: for all-correct or all-wrong, any choice = oracle -> frac = 1.0
        span = oracle - minimum
        is_tie = span < 1e-6
        per_query_frac = np.where(
            is_tie,
            1.0,
            (chosen - minimum) / np.maximum(span, 1e-6),
        )
        frac = per_query_frac.mean()

        # Per-query Spearman
        sp_corrs = []
        for i in range(len(indices)):
            if rews[i].std() > 1e-6:
                corr, _ = spearmanr(preds[i], rews[i])
                if not np.isnan(corr):
                    sp_corrs.append(corr)
        sp_mean = np.mean(sp_corrs) if sp_corrs else 0.0

        return frac, regret, sp_mean

    # Bootstrap
    bootstrap_frac = []
    bootstrap_regret = []
    bootstrap_spearman = []

    rng = np.random.RandomState(42)
    for _ in range(n_bootstrap):
        indices = rng.choice(n_queries, size=n_queries, replace=True)
        frac, regret, sp = compute_metrics(indices)
        bootstrap_frac.append(frac)
        bootstrap_regret.append(regret)
        bootstrap_spearman.append(sp)

    # Confidence intervals (2.5th and 97.5th percentiles)
    frac_ci = (float(np.percentile(bootstrap_frac, 2.5)), float(np.percentile(bootstrap_frac, 97.5)))
    regret_ci = (float(np.percentile(bootstrap_regret, 2.5)), float(np.percentile(bootstrap_regret, 97.5)))
    spearman_ci = (float(np.percentile(bootstrap_spearman, 2.5)), float(np.percentile(bootstrap_spearman, 97.5)))

    # Standard errors
    frac_se = float(np.std(bootstrap_frac))
    regret_se = float(np.std(bootstrap_regret))
    spearman_se = float(np.std(bootstrap_spearman))

    return StatisticalSignificance(
        frac_of_oracle_ci=frac_ci,
        regret_ci=regret_ci,
        spearman_ci=spearman_ci,
        frac_of_oracle_se=frac_se,
        regret_se=regret_se,
        spearman_se=spearman_se,
        n_samples=n_queries,
        n_effective=float(n_queries),  # Could compute effective sample size
    )


def compute_inference_benchmark(
    model: torch.nn.Module,
    sample_inputs: dict,
    n_warmup: int = 10,
    n_runs: int = 100,
    batch_sizes: list[int] | None = None,
) -> InferenceBenchmark:
    """Benchmark inference performance.

    Answers: "Is this fast enough for production?"

    Args:
        model: The PyTorch model
        sample_inputs: Dict of input tensors for forward pass
        n_warmup: Number of warmup runs
        n_runs: Number of timed runs
        batch_sizes: Batch sizes to test (default: [1])

    Returns:
        InferenceBenchmark with latency and throughput metrics
    """
    import time

    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype

    # Warmup
    model.eval()
    with torch.no_grad():
        for _ in range(n_warmup):
            _ = model(**sample_inputs)

    # Synchronize if GPU
    if device.type == "cuda":
        torch.cuda.synchronize()

    # Timed runs
    latencies = []
    with torch.no_grad():
        for _ in range(n_runs):
            start = time.perf_counter()
            _ = model(**sample_inputs)
            if device.type == "cuda":
                torch.cuda.synchronize()
            end = time.perf_counter()
            latencies.append((end - start) * 1000)  # ms

    latencies = np.array(latencies)

    # Model size
    param_count = sum(p.numel() for p in model.parameters())
    model_size_mb = param_count * 4 / (1024 * 1024)  # Assume float32

    # Peak memory (approximate)
    if device.type == "cuda":
        peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)
    else:
        peak_memory_mb = model_size_mb * 2  # Rough estimate

    batch_size = sample_inputs.get("query_features", sample_inputs.get("edge_mask", torch.zeros(1))).shape[0]

    return InferenceBenchmark(
        latency_mean_ms=float(latencies.mean()),
        latency_p50_ms=float(np.percentile(latencies, 50)),
        latency_p95_ms=float(np.percentile(latencies, 95)),
        latency_p99_ms=float(np.percentile(latencies, 99)),
        queries_per_second=float(1000 / latencies.mean() * batch_size),
        batch_size_tested=batch_size,
        model_size_mb=model_size_mb,
        peak_memory_mb=peak_memory_mb,
        device=str(device),
        dtype=str(dtype),
    )


@dataclass
class ModelPredictions:
    """GNN model predictions for evaluation."""

    logits: np.ndarray  # (N, num_llms) - raw routing logits
    uncertainties: np.ndarray | None  # (N, num_llms) or (N,) - predicted variances
    difficulty_preds: np.ndarray | None  # (N,) - predicted difficulty scores
    test_indices: np.ndarray  # Indices into the full dataset
    mean_logit: np.ndarray | None = None  # (N,) - mean logit per query (difficulty signal)
    llm_names: list[str] | None = None  # LLM names in order
    test_rewards: np.ndarray | None = None  # (N, num_llms) - rewards from cache


def load_model_and_predict(
    checkpoint_path: str | Path,
    device: str = "cpu",
) -> ModelPredictions:
    """Load GNN model and run inference on test set.

    This function loads a trained GNN router from a registry checkpoint and
    runs inference on the test set to get routing logits, difficulty predictions,
    and uncertainty estimates.

    Args:
        checkpoint_path: Path to registry checkpoint directory containing:
            - model.pth: Model weights
            - training_graph_cache.pkl: Cached graph with test indices
            - config.yaml: Model architecture config
        device: Device for inference ("cpu" or "cuda")

    Returns:
        ModelPredictions with logits, uncertainties, difficulty predictions
    """
    import pickle

    import yaml

    checkpoint_path = Path(checkpoint_path)

    # Load training graph cache
    cache_path = checkpoint_path / "training_graph_cache.pkl"
    if not cache_path.exists():
        raise FileNotFoundError(f"Training graph cache not found: {cache_path}")

    logger.info("Loading training graph cache from %s", cache_path)
    with open(cache_path, "rb") as f:
        cache = pickle.load(f)

    # Load config
    config_path = checkpoint_path / "config.yaml"
    if config_path.exists():
        with open(config_path) as f:
            raw_config = yaml.safe_load(f)
        # Normalize config keys to match model constructor
        config = {
            "hidden_features": raw_config.get("hidden_dim", 384),
            "num_gnn_layers": raw_config.get("num_gat_layers", 3),
            "num_attention_heads": raw_config.get("num_gat_heads", 4),
            "dropout": raw_config.get("dropout", 0.4),
            "use_difficulty_head": raw_config.get("use_difficulty_head", True),
            "hidden_multiplier": 2,  # Default for GAT
        }
    else:
        raw_config = {}
        # Default config matching the model architecture
        config = {
            "hidden_features": 384,
            "num_gnn_layers": 3,
            "dropout": 0.4,
            "use_difficulty_head": True,
            "num_attention_heads": 4,
            "hidden_multiplier": 2,
        }

    # Detect actual architecture from state_dict keys and tensor shapes
    # This prevents config file mismatches from causing load errors
    model_path = checkpoint_path / "model.pth"
    state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
    state_keys = set(state_dict.keys())

    # Detect contrastive encoder
    config["use_contrastive"] = any(k.startswith("contrastive_encoder.") for k in state_keys)
    # Detect uncertainty head
    config["use_uncertainty_head"] = any(k.startswith("uncertainty_proj.") for k in state_keys)

    # Infer hidden_features from state_dict tensor shapes
    # model_align.query_transform.weight has shape [hidden_features, query_dim]
    if "model_align.query_transform.weight" in state_dict:
        inferred_hidden = state_dict["model_align.query_transform.weight"].shape[0]
        inferred_query_dim = state_dict["model_align.query_transform.weight"].shape[1]
        config["hidden_features"] = inferred_hidden
        config["inferred_query_dim"] = inferred_query_dim
        logger.info("Inferred from checkpoint: hidden_features=%d, query_dim=%d",
                    inferred_hidden, inferred_query_dim)

    # Infer llm_dim from model_align.llm_transform.weight [hidden*2, llm_dim]
    if "model_align.llm_transform.weight" in state_dict:
        inferred_llm_dim = state_dict["model_align.llm_transform.weight"].shape[1]
        config["inferred_llm_dim"] = inferred_llm_dim
        logger.info("Inferred from checkpoint: llm_dim=%d", inferred_llm_dim)

    logger.info(
        "Config: hidden=%d, layers=%d, heads=%d, contrastive=%s, uncertainty=%s",
        config["hidden_features"], config["num_gnn_layers"],
        config["num_attention_heads"], config["use_contrastive"],
        config["use_uncertainty_head"],
    )

    # Extract cache contents
    query_embeddings = torch.from_numpy(cache["query_embeddings"]).float()
    task_embeddings = torch.from_numpy(cache["task_embeddings"]).float()
    llm_embeddings = torch.from_numpy(cache["llm_embeddings"]).float()
    combined_edge = torch.from_numpy(cache["combined_edge"]).float()
    source_nodes = torch.from_numpy(cache["source_nodes"]).long()
    dest_nodes = torch.from_numpy(cache["dest_nodes"]).long()
    test_mask = torch.from_numpy(cache["test_mask"]).bool()

    num_queries = cache["num_queries"]
    num_llms = cache["num_llms"]
    llm_order = cache.get("llm_order", [f"LLM_{i}" for i in range(num_llms)])

    logger.info(
        "Loaded cache: %d queries, %d LLMs, %d edges (%d test)",
        num_queries, num_llms, len(source_nodes), test_mask.sum().item(),
    )

    # Build edge_index: source is query node (0..num_queries-1),
    # dest is LLM node (num_queries..num_queries+num_llms-1)
    # The GNN operates on a combined node set: [queries | LLMs]
    edge_index = torch.stack([
        source_nodes,
        dest_nodes + num_queries,  # Offset LLM indices
    ])

    # Build edge features matching training-time construction.
    #
    # Training builds 6-dim features via _edge_features_with_pref():
    #   [wp, wc, wl, cost, perf, latency]  (preferences PREPENDED)
    #
    # Cache combined_edge can have two column orderings:
    #   - "cpl" (training native): [cost, perf, latency]
    #   - "pcl" (legacy cache):    [perf, cost, latency]
    #
    # Quality-only training uses wp=0, wc=0, wl=0 for edge features.
    n_edges = combined_edge.shape[0]
    cache_col_order = cache.get("col_order", "pcl")  # Default to legacy
    quality_only = raw_config.get("quality_only_training", False)

    if quality_only:
        wp, wc, wl = 0.0, 0.0, 0.0
    else:
        iw = raw_config.get("inference_weights", {})
        wp = float(iw.get("perf", 0.7))
        wc = float(iw.get("cost", 0.2))
        wl = float(iw.get("latency", 0.1))

    logger.info(
        "Edge features: col_order=%s, quality_only=%s, weights=(%.1f, %.1f, %.1f)",
        cache_col_order, quality_only, wp, wc, wl,
    )

    if combined_edge.shape[1] == 3:
        # Reorder to training column order [cost, perf, latency] if needed
        if cache_col_order == "pcl":
            # Cache has [perf, cost, latency] -> reorder to [cost, perf, latency]
            combined_edge_cpl = combined_edge[:, [1, 0, 2]]
        else:
            combined_edge_cpl = combined_edge

        # Prepend preference weights (matching _edge_features_with_pref)
        weights = torch.tensor(
            [wp, wc, wl], dtype=combined_edge_cpl.dtype,
        ).unsqueeze(0).expand(n_edges, -1)
        edge_weight = torch.cat([weights, combined_edge_cpl], dim=1)
    else:
        edge_weight = combined_edge

    # Initialize model
    from embergraphrouter.model.graph_nn import GATEncoderDecoderNet

    # Use inferred dimensions from checkpoint if available, otherwise from cache
    cache_query_dim = query_embeddings.shape[1]
    cache_llm_dim = llm_embeddings.shape[1]

    query_dim = config.get("inferred_query_dim", cache_query_dim)
    llm_dim = config.get("inferred_llm_dim", cache_llm_dim)

    # Handle dimension mismatch between cache and checkpoint
    if cache_query_dim != query_dim:
        logger.warning(
            "Cache query_dim=%d != checkpoint query_dim=%d, truncating/padding embeddings",
            cache_query_dim, query_dim,
        )
        if cache_query_dim > query_dim:
            # Truncate to match checkpoint (Matryoshka-style)
            query_embeddings = query_embeddings[:, :query_dim]
            task_embeddings = task_embeddings[:, :query_dim]
        else:
            # Pad with zeros
            pad = torch.zeros(query_embeddings.shape[0], query_dim - cache_query_dim)
            query_embeddings = torch.cat([query_embeddings, pad], dim=1)
            task_pad = torch.zeros(task_embeddings.shape[0], query_dim - cache_query_dim)
            task_embeddings = torch.cat([task_embeddings, task_pad], dim=1)

    if cache_llm_dim != llm_dim:
        logger.warning(
            "Cache llm_dim=%d != checkpoint llm_dim=%d, truncating/padding embeddings",
            cache_llm_dim, llm_dim,
        )
        if cache_llm_dim > llm_dim:
            llm_embeddings = llm_embeddings[:, :llm_dim]
        else:
            pad = torch.zeros(llm_embeddings.shape[0], llm_dim - cache_llm_dim)
            llm_embeddings = torch.cat([llm_embeddings, pad], dim=1)

    # Detect score_head from checkpoint state_dict (loaded earlier for arch detection)
    use_score_head = "score_head.weight" in state_dict

    model = GATEncoderDecoderNet(
        query_feature_dim=query_dim,
        llm_feature_dim=llm_dim,
        hidden_features=config["hidden_features"],
        in_edges=6,  # (perf, cost, latency, w_perf, w_cost, w_latency)
        dropout=config["dropout"],
        num_gnn_layers=config["num_gnn_layers"],
        use_difficulty_head=config["use_difficulty_head"],
        use_uncertainty_head=config["use_uncertainty_head"],
        num_attention_heads=config["num_attention_heads"],
        hidden_multiplier=config["hidden_multiplier"],
        use_contrastive=config["use_contrastive"],
        use_score_head=use_score_head,
    )

    # Load weights (reuse state_dict already loaded for architecture detection)
    # Use strict=False to ignore deprecated difficulty_head/uncertainty_head weights
    # from old checkpoints - these are now derived from mean logit instead.
    logger.info("Loading model weights from %s", model_path)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if unexpected:
        logger.info("Ignored deprecated weights from checkpoint: %s", unexpected)
    model.eval()
    model.to(device)

    # Move tensors to device
    query_embeddings = query_embeddings.to(device)
    task_embeddings = task_embeddings.to(device)
    llm_embeddings = llm_embeddings.to(device)
    edge_index = edge_index.to(device)
    edge_weight = edge_weight.to(device)
    test_mask = test_mask.to(device)

    # For inference, edge_can_see = all edges (full graph visibility)
    # edge_mask = test_mask (predict only on test edges)
    edge_can_see = torch.ones(n_edges, dtype=torch.bool, device=device)

    # Run inference
    # Note: return_difficulty and return_uncertainty are now deprecated.
    # Difficulty is computed from mean logit (0.31 correlation vs 0.05).
    logger.info("Running inference on test set...")
    with torch.no_grad():
        edge_logits = model(
            task_features=task_embeddings,
            query_features=query_embeddings,
            llm_features=llm_embeddings,
            edge_index=edge_index,
            edge_mask=test_mask,
            edge_can_see=edge_can_see,
            edge_weight=edge_weight,
        )

    # Reshape edge_logits from (num_test_edges,) to (num_test_queries, num_llms)
    # Test edges are organized by query: for each test query, we have num_llms edges
    test_source_nodes = source_nodes[test_mask.cpu()].numpy()
    test_dest_nodes = dest_nodes[test_mask.cpu()].numpy()

    # Extract performance values from cache.
    # Column layout depends on cache col_order:
    #   "pcl" (legacy): [perf=0, cost=1, latency=2]
    #   "cpl" (training native): [cost=0, perf=1, latency=2]
    perf_col = 0 if cache_col_order == "pcl" else 1
    test_perf = combined_edge[test_mask.cpu(), perf_col].numpy()

    # Get unique test queries and their order
    test_query_indices = np.unique(test_source_nodes)
    num_test_queries = len(test_query_indices)

    logger.info("Reshaping %d test edges to (%d queries, %d LLMs)",
                edge_logits.shape[0], num_test_queries, num_llms)

    # Build per-query logit matrix and rewards matrix
    logits_matrix = np.zeros((num_test_queries, num_llms))
    rewards_matrix = np.zeros((num_test_queries, num_llms))
    edge_logits_np = edge_logits.cpu().numpy()

    # Map test edges to matrix positions
    query_to_row = {q: i for i, q in enumerate(test_query_indices)}
    for i, (src, dst) in enumerate(zip(test_source_nodes, test_dest_nodes)):
        row = query_to_row[src]
        col = dst  # dest_nodes are already 0-indexed LLM indices
        logits_matrix[row, col] = edge_logits_np[i]
        rewards_matrix[row, col] = test_perf[i]

    # Compute mean logit per query - this is the key difficulty signal.
    # Mean logit provides 6x better correlation with difficulty (0.31 vs 0.05)
    # compared to the deprecated difficulty_head.
    mean_logit = logits_matrix.mean(axis=1)

    logger.info("Inference complete: logits=%s, mean_logit=%s, rewards=%s",
                logits_matrix.shape, mean_logit.shape, rewards_matrix.shape)

    return ModelPredictions(
        logits=logits_matrix,
        uncertainties=None,  # Deprecated: use mean_logit instead
        difficulty_preds=None,  # Deprecated: use mean_logit instead
        test_indices=test_query_indices,
        mean_logit=mean_logit,
        llm_names=llm_order,
        test_rewards=rewards_matrix,
    )


def run_evaluation(
    router_data_path: str,
    checkpoint_path: str | None = None,
    output_path: str | None = None,
    run_advanced: bool = False,
    run_benchmark: bool = False,
) -> EvaluationResults:
    """Run complete evaluation pipeline.

    Args:
        router_data_path: Path to router_data.csv or .parquet
        checkpoint_path: Path to model checkpoint (optional, for model evaluation)
        output_path: Path to save JSON results (optional)
        run_advanced: If True, run advanced analysis (error, cost, significance)
        run_benchmark: If True, run inference benchmarks (run on target hardware)

    Returns:
        EvaluationResults with all metrics
    """
    # Ground truth analysis
    analyzer = GroundTruthDifficultyAnalyzer(router_data_path)
    analyzer.print_summary()

    # Get reward matrix for routing evaluation
    rewards, query_ids, llm_names = analyzer.get_reward_matrix()
    n_queries, n_llms = rewards.shape

    # Build results
    results = EvaluationResults(
        n_queries=n_queries,
        n_llms=n_llms,
        difficulty_distribution=analyzer.get_difficulty_distribution(),
        task_difficulty=analyzer.get_task_difficulty(),
    )

    # If checkpoint provided, run model evaluation
    if checkpoint_path:
        logger.info("Loading model from %s", checkpoint_path)
        try:
            model_preds = load_model_and_predict(checkpoint_path)

            # Use rewards from cache (perfectly aligned with model predictions)
            # This avoids index misalignment between cache and data file.
            if model_preds.test_rewards is not None:
                test_rewards = model_preds.test_rewards
                logger.info("Using rewards from cache (%d queries)", len(test_rewards))
            else:
                # Fallback for old caches without embedded rewards
                test_rewards = rewards[model_preds.test_indices]
                logger.warning("Cache has no rewards; using data file (may be misaligned)")
            test_difficulty = (test_rewards >= 0.5).sum(axis=1)

            # Compute overall routing metrics
            results.overall_routing = compute_routing_metrics(
                model_preds.logits, test_rewards
            )
            logger.info("Overall frac_oracle: %.3f", results.overall_routing.frac_of_oracle)

            # Compute routing metrics by difficulty band
            results.routing_by_band = compute_routing_by_difficulty(
                model_preds.logits, test_rewards, test_difficulty, n_llms
            )

            # Compute calibration metrics (includes "knows when it doesn't know" tests)
            results.calibration = compute_calibration_metrics(
                model_preds.logits, test_rewards
            )
            logger.info("ECE: %.3f", results.calibration.ece)

            # Print comprehensive "Knows When It Doesn't Know" analysis
            print("\n--- Entropy Analysis: Does the Router Know When It Doesn't Know? ---")
            cal = results.calibration
            print(f"  Entropy distribution:")
            print(f"    avg_entropy (normalized): {cal.avg_entropy:.3f} (1.0 = uniform)")
            print(f"    entropy_std:              {cal.entropy_std:.3f}")
            print(f"    % queries with entropy > 0.9: {cal.pct_high_entropy:.1f}%")
            print()
            print(f"  Entropy correlations:")
            print(f"    entropy ↔ difficulty:  {cal.entropy_difficulty_corr:+.3f}")
            print(f"    entropy ↔ error:       {cal.entropy_error_corr:+.3f} (want POSITIVE)")
            print(f"    entropy ↔ quality:     {cal.entropy_quality_corr:+.3f} (want NEGATIVE)")
            print()
            # Interpret the entropy results
            if cal.pct_high_entropy > 90:
                print("  ⚠ ENTROPY NOT USEFUL: >90% queries have near-uniform entropy.")
                print("    Softmax compresses small logit differences into uniform probabilities.")
            elif cal.entropy_error_corr > 0.1 and cal.entropy_quality_corr < -0.1:
                print("  ✓ ENTROPY IS USEFUL: Model knows when it doesn't know!")
                print(f"    High entropy predicts errors (r={cal.entropy_error_corr:+.2f})")
                print(f"    High entropy predicts low quality (r={cal.entropy_quality_corr:+.2f})")
            else:
                print("  ⚠ ENTROPY WEAKLY USEFUL: Correlations are weak.")

            # Raw logit uncertainty (bypasses softmax compression)
            print()
            print("  Raw logit uncertainty (bypassing softmax):")
            print(f"    avg_logit_std:    {cal.avg_logit_std:.3f}")
            print(f"    avg_logit_margin: {cal.avg_logit_margin:.3f} (max - 2nd max)")
            print()
            print(f"  Raw logit correlations:")
            print(f"    logit_std ↔ error:   {cal.logit_std_error_corr:+.3f} (want NEGATIVE)")
            print(f"    logit_std ↔ quality: {cal.logit_std_quality_corr:+.3f} (want POSITIVE)")
            print(f"    margin ↔ error:      {cal.logit_margin_error_corr:+.3f} (want NEGATIVE)")
            print(f"    margin ↔ quality:    {cal.logit_margin_quality_corr:+.3f} (want POSITIVE)")
            print()

            # Interpret raw logit uncertainty
            std_useful = cal.logit_std_error_corr < -0.05 or cal.logit_std_quality_corr > 0.05
            margin_useful = cal.logit_margin_error_corr < -0.05 or cal.logit_margin_quality_corr > 0.05

            if std_useful or margin_useful:
                print("  ✓ RAW LOGIT UNCERTAINTY IS USEFUL:")
                if std_useful:
                    print(f"    logit_std predicts routing {'errors' if cal.logit_std_error_corr < -0.05 else 'quality'}")
                if margin_useful:
                    print(f"    logit_margin predicts routing {'errors' if cal.logit_margin_error_corr < -0.05 else 'quality'}")
            else:
                print("  ⚠ RAW LOGIT UNCERTAINTY NOT USEFUL: weak correlations.")

            # Compute difficulty prediction correlation if available
            if model_preds.difficulty_preds is not None:
                results.difficulty_pred_corr = np.corrcoef(
                    model_preds.difficulty_preds, test_difficulty
                )[0, 1]
                logger.info(
                    "Difficulty prediction correlation: %.3f",
                    results.difficulty_pred_corr,
                )

            # Compute uncertainty-error correlation if available
            if model_preds.uncertainties is not None:
                # Prediction error: 1 if wrong top-1, 0 if correct
                pred_idx = model_preds.logits.argmax(axis=1)
                true_idx = test_rewards.argmax(axis=1)
                errors = (pred_idx != true_idx).astype(float)

                # Uncertainty can be per-edge or per-query
                if model_preds.uncertainties.ndim == 1:
                    uncertainty = model_preds.uncertainties
                else:
                    uncertainty = model_preds.uncertainties.mean(axis=1)

                results.uncertainty_error_corr = np.corrcoef(uncertainty, errors)[0, 1]
                logger.info(
                    "Uncertainty-error correlation: %.3f",
                    results.uncertainty_error_corr,
                )

            # Compute edge case calibration (impossible/trivial queries)
            # Find cheapest model index (GPT-5 Nano is typically index 2)
            cheap_idx = 2  # GPT-5 Nano
            if "GPT-5 Nano" in llm_names:
                cheap_idx = llm_names.index("GPT-5 Nano")

            results.edge_case_calibration = compute_edge_case_calibration(
                model_preds.logits,
                test_rewards,
                n_llms=n_llms,
                cheap_model_idx=cheap_idx,
            )

            ecc = results.edge_case_calibration
            logger.info("Edge case calibration:")
            logger.info(
                "  Impossible: %d queries (%.1f%%), avg_entropy=%.3f, overconf=%.3f",
                ecc.n_impossible, ecc.pct_impossible,
                ecc.impossible_avg_entropy, ecc.impossible_overconfidence,
            )
            logger.info(
                "  Trivial: %d queries (%.1f%%), avg_entropy=%.3f",
                ecc.n_trivial, ecc.pct_trivial, ecc.trivial_avg_entropy,
            )
            logger.info(
                "  Cost optimization: %.1f%% safe for cheap routing, "
                "potential savings=%.1f%%",
                100 * ecc.cost_opt_safe_cheap_pct,
                100 * ecc.cost_opt_potential_savings,
            )

            # Print routing by difficulty band
            print("\n--- Routing Metrics by Difficulty Band ---")
            for band_name in ["trivial", "easy", "medium", "hard", "very_hard", "impossible"]:
                if band_name not in results.routing_by_band:
                    continue
                metrics = results.routing_by_band[band_name]
                print(
                    f"  {band_name:12s}: frac_oracle={metrics.frac_of_oracle:.3f}, "
                    f"regret={metrics.regret:.3f}, top1={metrics.top1_acc:.3f}"
                )

            # Compute difficulty detection metrics (key insight: mean logit is the signal)
            if model_preds.mean_logit is not None and model_preds.llm_names is not None:
                results.difficulty_detection = compute_difficulty_detection_metrics(
                    model_preds.mean_logit,
                    test_rewards,
                    model_preds.llm_names,
                )

                dd = results.difficulty_detection
                print("\n--- Difficulty Detection (Mean Logit Analysis) ---")
                print("Mean logit by difficulty band:")
                for band in ["trivial", "medium", "hard", "very_hard", "impossible"]:
                    if band in dd.mean_logit_by_band:
                        ml = dd.mean_logit_by_band[band]
                        std = dd.mean_logit_std_by_band.get(band, 0)
                        print(f"  {band:12s}: {ml:+.3f} +/- {std:.3f}")

                print(f"\nTrivial vs Impossible separation:")
                print(f"  AUC: {dd.trivial_vs_impossible_auc:.3f}")
                print(f"  Optimal threshold: {dd.optimal_threshold:.3f}")
                print(f"  Separation accuracy: {dd.separation_accuracy*100:.1f}%")

                print(f"\nCorrelation(mean_logit, n_correct): {dd.mean_logit_difficulty_corr:.3f}")

                print(f"\nRecommended production thresholds:")
                print(f"  Trivial (route to cheap): mean_logit > {dd.recommended_trivial_threshold:.2f}")
                print(f"  Hard (use fallback LLMs): mean_logit < {dd.recommended_hard_threshold:.2f}")

                print(f"\nBest LLMs for hard queries: {dd.best_llms_on_hard}")
                logger.info(
                    "Difficulty detection: AUC=%.3f, corr=%.3f",
                    dd.trivial_vs_impossible_auc,
                    dd.mean_logit_difficulty_corr,
                )

            # Compute ranking discrimination metrics
            # Get task IDs for test queries if available
            test_task_ids = None
            if hasattr(analyzer, 'query_stats') and 'task_id' in analyzer.query_stats.columns:
                # Map test query indices to task IDs
                query_to_task = dict(zip(
                    range(len(analyzer.query_stats)),
                    analyzer.query_stats['task_id'].values
                ))
                test_task_ids = np.array([
                    query_to_task.get(idx, 'unknown')
                    for idx in model_preds.test_indices
                ])

            results.ranking_discrimination = compute_ranking_discrimination(
                model_preds.logits,
                test_rewards,
                model_preds.llm_names or llm_names,
                task_ids=test_task_ids,
            )

            rd = results.ranking_discrimination
            print("\n--- Ranking Discrimination: Are the Router's Rankings Meaningful? ---")
            print(f"  Per-query ranking correlation:")
            print(f"    avg Spearman:    {rd.avg_spearman_corr:+.3f}")
            print(f"    avg Kendall τ:   {rd.avg_kendall_tau:+.3f}")
            print(f"    % positive corr: {rd.pct_positive_corr:.1f}%")
            print()
            print(f"  Per-LLM discrimination (high score → high reward?):")
            for llm_name, disc in sorted(rd.per_llm_discrimination.items(), key=lambda x: -x[1]):
                indicator = "✓" if disc > 0.1 else "~" if disc > 0 else "✗"
                print(f"    {indicator} {llm_name:25s}: {disc:+.3f}")
            print(f"    Average:                       {rd.avg_llm_discrimination:+.3f}")
            print()
            print(f"  Overall score-reward calibration:")
            print(f"    corr(all_scores, all_rewards): {rd.score_reward_corr:+.3f}")
            print(f"    P(correct | top 25% score):    {rd.high_score_success_rate:.1%}")
            print(f"    P(wrong | bottom 25% score):   {rd.low_score_failure_rate:.1%}")

            if test_task_ids is not None and rd.task_affinity_captured:
                print()
                print(f"  Task-specific affinity capture:")
                print(f"    avg task affinity corr: {rd.avg_task_affinity:+.3f}")
                if rd.tasks_with_strong_affinity:
                    task_names = [str(t) for t in rd.tasks_with_strong_affinity[:5]]
                    print(f"    tasks with strong affinity (>0.5): {task_names}")

            # Interpret discrimination results
            print()
            if rd.avg_spearman_corr > 0.3 and rd.avg_llm_discrimination > 0.1:
                print("  ✓ RANKINGS ARE DISCRIMINATIVE: Router predictions are meaningful!")
            elif rd.avg_spearman_corr > 0.1 or rd.avg_llm_discrimination > 0.05:
                print("  ~ RANKINGS ARE WEAKLY DISCRIMINATIVE: Some signal, but noisy.")
            else:
                print("  ✗ RANKINGS NOT DISCRIMINATIVE: Router predictions are not meaningful.")

            logger.info(
                "Ranking discrimination: Spearman=%.3f, LLM_disc=%.3f, score_reward=%.3f",
                rd.avg_spearman_corr, rd.avg_llm_discrimination, rd.score_reward_corr,
            )

            # Advanced analysis (opt-in for production use)
            if run_advanced:
                logger.info("Running advanced analysis...")

                # Error analysis
                results.error_analysis = compute_error_analysis(
                    model_preds.logits,
                    test_rewards,
                    model_preds.llm_names or llm_names,
                    task_ids=test_task_ids,
                )
                ea = results.error_analysis
                print("\n--- Error Analysis ---")
                print(f"  Errors: {ea.n_errors}/{len(model_preds.logits)} "
                      f"({ea.error_rate:.1%})")
                print(f"  Avg regret on errors: {ea.avg_error_regret:.3f}")
                if ea.most_confused_pairs:
                    print("  Most confused pairs:")
                    for pred, actual, cnt in ea.most_confused_pairs[:3]:
                        print(f"    {pred} -> {actual}: {cnt} times")

                # Cost-quality analysis
                results.cost_quality = compute_cost_quality_analysis(
                    model_preds.logits,
                    test_rewards,
                    model_preds.llm_names or llm_names,
                )
                cq = results.cost_quality
                print("\n--- Cost-Quality Analysis ---")
                print(f"  Router cost:    {cq.router_routing_cost:.2f}")
                print(f"  Oracle cost:    {cq.oracle_routing_cost:.2f}")
                print(f"  Random cost:    {cq.random_routing_cost:.2f}")
                print(f"  Cost savings vs random: {cq.cost_savings_vs_random:.1%}")

                # Statistical significance
                results.statistical_significance = compute_statistical_significance(
                    model_preds.logits,
                    test_rewards,
                )
                ss = results.statistical_significance
                print("\n--- Statistical Significance (95% CI) ---")
                print(f"  frac_of_oracle: {ss.frac_of_oracle_ci[0]:.3f} - "
                      f"{ss.frac_of_oracle_ci[1]:.3f}")
                print(f"  regret:         {ss.regret_ci[0]:.3f} - "
                      f"{ss.regret_ci[1]:.3f}")
                print(f"  spearman:       {ss.spearman_ci[0]:.3f} - "
                      f"{ss.spearman_ci[1]:.3f}")

            # Inference benchmark (opt-in, run on target hardware like B200)
            if run_benchmark:
                logger.info("Running inference benchmarks...")
                # Need to reload model for benchmarking
                from embergraphrouter.model.graph_nn import GATEncoderDecoderNet
                model_path = Path(checkpoint_path) / "model.pth"
                cache_path = Path(checkpoint_path) / "training_graph_cache.pkl"

                import pickle
                with open(cache_path, "rb") as f:
                    cache = pickle.load(f)

                # Rebuild sample inputs for benchmarking
                query_emb = torch.from_numpy(cache["query_embeddings"]).float()
                llm_emb = torch.from_numpy(cache["llm_embeddings"]).float()
                task_emb = torch.from_numpy(cache["task_embeddings"]).float()
                src = torch.from_numpy(cache["source_nodes"]).long()
                dst = torch.from_numpy(cache["dest_nodes"]).long()
                edge_feat = torch.from_numpy(cache["combined_edge"]).float()
                test_mask = torch.from_numpy(cache["test_mask"]).bool()

                n_edges = edge_feat.shape[0]
                if edge_feat.shape[1] == 3:
                    w = torch.ones(n_edges, 3) * torch.tensor([0.7, 0.2, 0.1])
                    edge_feat = torch.cat([edge_feat, w], dim=1)

                edge_idx = torch.stack([src, dst + cache["num_queries"]])
                can_see = torch.ones(n_edges, dtype=torch.bool)

                sample_inputs = {
                    "task_features": task_emb,
                    "query_features": query_emb,
                    "llm_features": llm_emb,
                    "edge_index": edge_idx,
                    "edge_mask": test_mask,
                    "edge_can_see": can_see,
                    "edge_weight": edge_feat,
                }

                # Load model again for clean benchmarking
                state = torch.load(model_path, map_location="cpu", weights_only=True)
                model = GATEncoderDecoderNet(
                    query_feature_dim=query_emb.shape[1],
                    llm_feature_dim=llm_emb.shape[1],
                    hidden_features=384,
                    in_edges=6,
                    dropout=0.0,  # No dropout for inference
                    num_gnn_layers=3,
                    use_difficulty_head=False,
                    use_uncertainty_head=False,
                    num_attention_heads=4,
                    hidden_multiplier=2,
                    use_contrastive=False,
                )
                model.load_state_dict(state, strict=False)
                model.eval()

                results.inference_benchmark = compute_inference_benchmark(
                    model, sample_inputs
                )
                ib = results.inference_benchmark
                print("\n--- Inference Benchmark ---")
                print(f"  Latency: {ib.latency_mean_ms:.2f}ms mean, "
                      f"{ib.latency_p95_ms:.2f}ms p95")
                print(f"  Throughput: {ib.queries_per_second:.1f} queries/sec")
                print(f"  Model size: {ib.model_size_mb:.1f} MB")
                print(f"  Device: {ib.device}")

        except NotImplementedError as e:
            logger.warning("Model evaluation skipped: %s", e)

    # Save results
    if output_path:
        output_file = Path(output_path)
        output_file.parent.mkdir(parents=True, exist_ok=True)
        with open(output_file, "w") as f:
            json.dump(results.to_dict(), f, indent=2, default=_json_default)
        logger.info("Results saved to %s", output_file)

    return results


def print_scorecard(results: EvaluationResults, name: str = "Router") -> dict[str, float]:
    """Print a concise scorecard of key metrics.

    Returns dict of metrics for comparison.
    """
    metrics = {}

    print(f"\n{'=' * 60}")
    print(f"  SCORECARD: {name}")
    print(f"{'=' * 60}")

    if results.overall_routing:
        metrics["frac_of_oracle"] = results.overall_routing.frac_of_oracle
        metrics["regret"] = results.overall_routing.regret
        metrics["top1_acc"] = results.overall_routing.top1_acc
        print(f"  Routing Quality:")
        print(f"    frac_of_oracle:  {metrics['frac_of_oracle']:.3f}")
        print(f"    regret:          {metrics['regret']:.3f}")

    if results.difficulty_detection:
        dd = results.difficulty_detection
        metrics["difficulty_corr"] = dd.mean_logit_difficulty_corr
        metrics["difficulty_auc"] = dd.trivial_vs_impossible_auc
        print(f"  Difficulty Detection:")
        print(f"    mean_logit_corr: {metrics['difficulty_corr']:.3f}")
        print(f"    trivial_vs_imp:  {metrics['difficulty_auc']:.3f} AUC")

    if results.calibration:
        cal = results.calibration
        metrics["entropy_error_corr"] = cal.entropy_error_corr
        metrics["logit_std_quality_corr"] = cal.logit_std_quality_corr
        print(f"  Uncertainty Signal:")
        print(f"    entropy→error:   {metrics['entropy_error_corr']:+.3f}")
        print(f"    logit_std→qual:  {metrics['logit_std_quality_corr']:+.3f}")

    if results.ranking_discrimination:
        rd = results.ranking_discrimination
        metrics["spearman_corr"] = rd.avg_spearman_corr
        metrics["llm_discrimination"] = rd.avg_llm_discrimination
        metrics["task_affinity"] = rd.avg_task_affinity
        metrics["score_reward_corr"] = rd.score_reward_corr
        print(f"  Ranking Discrimination:")
        print(f"    spearman:        {metrics['spearman_corr']:+.3f}")
        print(f"    llm_disc:        {metrics['llm_discrimination']:+.3f}")
        print(f"    task_affinity:   {metrics['task_affinity']:+.3f}")

    print(f"{'=' * 60}\n")
    return metrics


def compare_routers(
    results_a: EvaluationResults,
    results_b: EvaluationResults,
    name_a: str = "Router A",
    name_b: str = "Router B",
) -> None:
    """Compare two routers and print side-by-side analysis.

    Prints a clear comparison showing which router is better on each metric.
    """
    metrics_a = print_scorecard(results_a, name_a)
    metrics_b = print_scorecard(results_b, name_b)

    print("\n" + "=" * 70)
    print("  HEAD-TO-HEAD COMPARISON")
    print("=" * 70)

    # Define metrics where higher is better vs lower is better
    higher_better = {
        "frac_of_oracle": "Routing Quality",
        "top1_acc": "Top-1 Accuracy",
        "difficulty_corr": "Difficulty Detection",
        "difficulty_auc": "Difficulty AUC",
        "spearman_corr": "Ranking Correlation",
        "llm_discrimination": "LLM Discrimination",
        "task_affinity": "Task Affinity",
        "score_reward_corr": "Score-Reward Calibration",
        "logit_std_quality_corr": "Uncertainty→Quality",
    }
    lower_better = {
        "regret": "Regret",
    }
    # For entropy_error_corr, higher (more positive) is better
    higher_better["entropy_error_corr"] = "Uncertainty→Error"

    a_wins = 0
    b_wins = 0
    ties = 0

    print(f"\n  {'Metric':<25} {'A':>10} {'B':>10} {'Winner':>12}")
    print(f"  {'-' * 25} {'-' * 10} {'-' * 10} {'-' * 12}")

    for metric, label in {**higher_better, **lower_better}.items():
        if metric not in metrics_a or metric not in metrics_b:
            continue

        val_a = metrics_a[metric]
        val_b = metrics_b[metric]

        # Determine winner
        threshold = 0.01  # Minimum difference to declare a winner
        if metric in higher_better:
            if val_a > val_b + threshold:
                winner = name_a
                a_wins += 1
            elif val_b > val_a + threshold:
                winner = name_b
                b_wins += 1
            else:
                winner = "TIE"
                ties += 1
        else:  # lower is better
            if val_a < val_b - threshold:
                winner = name_a
                a_wins += 1
            elif val_b < val_a - threshold:
                winner = name_b
                b_wins += 1
            else:
                winner = "TIE"
                ties += 1

        # Format winner with indicator
        if winner == name_a:
            winner_str = f"← {name_a}"
        elif winner == name_b:
            winner_str = f"{name_b} →"
        else:
            winner_str = "TIE"

        print(f"  {label:<25} {val_a:>10.3f} {val_b:>10.3f} {winner_str:>12}")

    print(f"\n  {'-' * 60}")
    print(f"  VERDICT: {name_a} wins {a_wins}, {name_b} wins {b_wins}, {ties} ties")

    if a_wins > b_wins:
        print(f"\n  >>> {name_a} IS BETTER OVERALL <<<")
    elif b_wins > a_wins:
        print(f"\n  >>> {name_b} IS BETTER OVERALL <<<")
    else:
        print(f"\n  >>> ROUTERS ARE COMPARABLE <<<")

    print("=" * 70 + "\n")


def run_comparison(
    router_data_path: str,
    checkpoint_a: str,
    checkpoint_b: str,
    name_a: str = "Router A",
    name_b: str = "Router B",
    run_advanced: bool = False,
    run_benchmark: bool = False,
) -> tuple[EvaluationResults, EvaluationResults]:
    """Run evaluation on two routers and compare them.

    Args:
        router_data_path: Path to router_data.csv
        checkpoint_a: Path to first checkpoint
        checkpoint_b: Path to second checkpoint
        name_a: Display name for first router
        name_b: Display name for second router
        run_advanced: If True, run advanced analysis
        run_benchmark: If True, run inference benchmarks

    Returns:
        Tuple of (results_a, results_b)
    """
    print(f"\n{'#' * 70}")
    print(f"  EVALUATING: {name_a}")
    print(f"{'#' * 70}")
    results_a = run_evaluation(
        router_data_path, checkpoint_a,
        run_advanced=run_advanced, run_benchmark=run_benchmark,
    )

    print(f"\n{'#' * 70}")
    print(f"  EVALUATING: {name_b}")
    print(f"{'#' * 70}")
    results_b = run_evaluation(
        router_data_path, checkpoint_b,
        run_advanced=run_advanced, run_benchmark=run_benchmark,
    )

    compare_routers(results_a, results_b, name_a, name_b)

    return results_a, results_b


def save_to_registry(results: EvaluationResults, checkpoint_path: Path) -> Path:
    """Save evaluation results to the registry alongside the model.

    Creates an eval_metrics.json file in the checkpoint directory with
    all evaluation metrics. This allows tracking model quality over time.

    Args:
        results: Evaluation results to save
        checkpoint_path: Path to the checkpoint directory

    Returns:
        Path to the saved metrics file
    """
    metrics_path = checkpoint_path / "eval_metrics.json"

    # Load existing metrics if present (to preserve history)
    history = []
    if metrics_path.exists():
        try:
            with open(metrics_path) as f:
                existing = json.load(f)
                if isinstance(existing, list):
                    history = existing
                else:
                    # Old format was single dict, convert to list
                    history = [existing]
        except json.JSONDecodeError:
            pass

    # Add timestamp and current results
    from datetime import datetime
    entry = {
        "timestamp": datetime.now().isoformat(),
        "metrics": results.to_dict(),
    }
    history.append(entry)

    # Save updated history
    with open(metrics_path, "w") as f:
        json.dump(history, f, indent=2)

    logger.info("Saved evaluation metrics to %s", metrics_path)

    # Also save a summary scorecard as a simple text file
    summary_path = checkpoint_path / "eval_summary.txt"
    with open(summary_path, "w") as f:
        f.write(f"Evaluation Summary - {datetime.now().strftime('%Y-%m-%d %H:%M')}\n")
        f.write("=" * 50 + "\n\n")

        if results.overall_routing:
            f.write("Routing Quality:\n")
            f.write(f"  frac_of_oracle: {results.overall_routing.frac_of_oracle:.3f}\n")
            f.write(f"  regret:         {results.overall_routing.regret:.3f}\n\n")

        if results.difficulty_detection:
            dd = results.difficulty_detection
            f.write("Difficulty Detection:\n")
            f.write(f"  mean_logit_corr: {dd.mean_logit_difficulty_corr:.3f}\n")
            f.write(f"  AUC:             {dd.trivial_vs_impossible_auc:.3f}\n\n")

        if results.ranking_discrimination:
            rd = results.ranking_discrimination
            f.write("Ranking Discrimination:\n")
            f.write(f"  spearman:        {rd.avg_spearman_corr:.3f}\n")
            f.write(f"  llm_disc:        {rd.avg_llm_discrimination:.3f}\n")
            f.write(f"  task_affinity:   {rd.avg_task_affinity:.3f}\n\n")

        if results.calibration:
            cal = results.calibration
            f.write("Uncertainty Signal:\n")
            f.write(f"  entropy→error:   {cal.entropy_error_corr:+.3f}\n")
            f.write(f"  logit_std→qual:  {cal.logit_std_quality_corr:+.3f}\n")

    logger.info("Saved evaluation summary to %s", summary_path)
    return metrics_path


def main():
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Difficulty-aware evaluation for GNN router",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Evaluate a single checkpoint
  python eval_difficulty_aware.py --checkpoint registry/eb1-gnn-preview-12-19

  # Compare two checkpoints
  python eval_difficulty_aware.py --checkpoint registry/model-v1 --compare registry/model-v2

  # Run with advanced analysis (error, cost, significance)
  python eval_difficulty_aware.py --checkpoint registry/my-model --advanced

  # Run inference benchmarks on B200 (run this on the target hardware)
  python eval_difficulty_aware.py --checkpoint registry/my-model --run-benchmark

  # Ground truth analysis only (no model)
  python eval_difficulty_aware.py --data-only
""",
    )
    parser.add_argument(
        "--data",
        type=str,
        default="data/router_data.csv",
        help="Path to router_data.csv or .parquet",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        default=None,
        help="Path to model checkpoint directory",
    )
    parser.add_argument(
        "--compare",
        type=str,
        default=None,
        help="Path to second checkpoint for comparison",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Path to save JSON results",
    )
    parser.add_argument(
        "--save-to-registry",
        action="store_true",
        default=True,
        help="Save results to checkpoint directory (default: True)",
    )
    parser.add_argument(
        "--no-save-to-registry",
        action="store_true",
        help="Disable saving results to checkpoint directory",
    )
    parser.add_argument(
        "--data-only",
        action="store_true",
        help="Run only ground truth analysis (no model)",
    )
    parser.add_argument(
        "--scorecard",
        action="store_true",
        help="Print concise scorecard at end",
    )
    parser.add_argument(
        "--advanced",
        action="store_true",
        help="Run advanced analysis (error, cost, significance)",
    )
    parser.add_argument(
        "--run-benchmark",
        action="store_true",
        help="Run inference benchmarks (should be run on target hardware, e.g. B200)",
    )

    args = parser.parse_args()

    # Resolve paths relative to embergraphrouter directory
    script_dir = Path(__file__).parent.parent
    data_path = script_dir / args.data if not Path(args.data).is_absolute() else Path(args.data)

    checkpoint = None
    if args.checkpoint and not args.data_only:
        checkpoint = script_dir / args.checkpoint if not Path(args.checkpoint).is_absolute() else Path(args.checkpoint)

    output = None
    if args.output:
        output = script_dir / args.output if not Path(args.output).is_absolute() else Path(args.output)

    # Comparison mode
    if args.compare:
        compare_path = script_dir / args.compare if not Path(args.compare).is_absolute() else Path(args.compare)

        if not checkpoint:
            parser.error("--compare requires --checkpoint to specify the first model")

        name_a = checkpoint.name
        name_b = compare_path.name

        results_a, results_b = run_comparison(
            router_data_path=str(data_path),
            checkpoint_a=str(checkpoint),
            checkpoint_b=str(compare_path),
            name_a=name_a,
            name_b=name_b,
            run_advanced=args.advanced,
            run_benchmark=args.run_benchmark,
        )

        # Save to both registries if enabled
        if args.save_to_registry and not args.no_save_to_registry:
            save_to_registry(results_a, checkpoint)
            save_to_registry(results_b, compare_path)

        return

    # Single evaluation mode
    results = run_evaluation(
        router_data_path=str(data_path),
        checkpoint_path=str(checkpoint) if checkpoint else None,
        output_path=str(output) if output else None,
        run_advanced=args.advanced,
        run_benchmark=args.run_benchmark,
    )

    # Save to registry if checkpoint provided and not disabled
    if checkpoint and args.save_to_registry and not args.no_save_to_registry:
        save_to_registry(results, checkpoint)

    # Print scorecard if requested
    if args.scorecard and checkpoint:
        print_scorecard(results, checkpoint.name)


if __name__ == "__main__":
    main()
