#!/usr/bin/env python3
"""Comprehensive baseline router training and difficulty-aware evaluation.

This script performs thorough training and evaluation of all three baseline
router architectures across multiple embedding dimensions, with detailed
difficulty-aware metrics and error analysis.

Key features:
1. Trains all baselines (Two Tower, Multi-Head MLP, GNN) with proper checkpointing
2. Supports multiple embedding dimensions (256, 512, 1024)
3. Computes difficulty-aware metrics stratified by query hardness
4. Analyzes error patterns and failure modes
5. Produces comprehensive comparison reports

Usage:
    python -m embergraphrouter.model.baselines.comprehensive_eval \
        --data_path embergraphrouter/data/router_data_8llm_qwen3.parquet \
        --output_dir results/baselines_comprehensive \
        --embedding_dims 256 512 1024 \
        --num_epochs 100
"""

from __future__ import annotations

import argparse
import ast
import json
import logging
import os
import pickle
import time
from collections import defaultdict
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy import stats
from torch.utils.data import DataLoader, Dataset, Subset

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


# ==============================================================================
# Data Classes for Results
# ==============================================================================

DIFFICULTY_BANDS = {
    "impossible": (0, 0),
    "very_hard": (1, 2),
    "hard": (3, 4),
    "medium": (5, 6),
    "easy": (7, 7),
    "trivial": (8, float("inf")),
}


@dataclass
class DifficultyMetrics:
    """Metrics computed per difficulty band."""

    band_name: str
    n_queries: int
    pct_of_total: float
    top1_acc: float
    frac_of_oracle: float
    avg_regret: float
    avg_effect: float


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

    ece: float  # Expected Calibration Error
    entropy_error_corr: float  # Correlation(entropy, error)
    entropy_quality_corr: float  # Correlation(entropy, frac_oracle)
    logit_margin_error_corr: float  # Correlation(margin, error)
    avg_entropy: float
    avg_max_prob: float


@dataclass
class ErrorAnalysis:
    """Detailed error analysis."""

    n_errors: int
    error_rate: float
    avg_error_regret: float
    errors_by_band: dict[str, int]
    error_rate_by_band: dict[str, float]
    worst_error_regrets: list[float]
    confusion_matrix: dict[str, dict[str, int]]
    most_confused_pairs: list[tuple[str, str, int]]


@dataclass
class RankingMetrics:
    """Ranking quality metrics."""

    avg_spearman_corr: float
    avg_kendall_tau: float
    pct_positive_corr: float
    ndcg_at_1: float
    ndcg_at_3: float


@dataclass
class BaselineResult:
    """Complete results for a single baseline model."""

    name: str
    embedding_dim: int
    params: int
    training_time: float
    best_epoch: int

    # Core metrics
    test_frac_oracle: float
    test_top1_acc: float
    test_regret: float

    # Training history
    train_loss_history: list[float]
    val_loss_history: list[float]
    val_frac_oracle_history: list[float]

    # Difficulty-aware metrics
    difficulty_metrics: list[DifficultyMetrics]

    # Calibration
    calibration: CalibrationMetrics

    # Error analysis
    errors: ErrorAnalysis

    # Ranking quality
    ranking: RankingMetrics

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        result = {
            "name": self.name,
            "embedding_dim": self.embedding_dim,
            "params": self.params,
            "training_time": self.training_time,
            "best_epoch": self.best_epoch,
            "test_frac_oracle": self.test_frac_oracle,
            "test_top1_acc": self.test_top1_acc,
            "test_regret": self.test_regret,
            "train_loss_history": self.train_loss_history,
            "val_loss_history": self.val_loss_history,
            "val_frac_oracle_history": self.val_frac_oracle_history,
            "difficulty_metrics": [asdict(m) for m in self.difficulty_metrics],
            "calibration": asdict(self.calibration),
            "errors": {
                **asdict(self.errors),
                "most_confused_pairs": [
                    {"predicted": p, "actual": a, "count": c}
                    for p, a, c in self.errors.most_confused_pairs
                ],
            },
            "ranking": asdict(self.ranking),
        }
        return result


# ==============================================================================
# Dataset
# ==============================================================================


def parse_embedding(emb_val: str | list | np.ndarray) -> np.ndarray:
    """Parse embedding from string or array."""
    if isinstance(emb_val, str):
        return np.array(ast.literal_eval(emb_val), dtype=np.float32)
    elif isinstance(emb_val, (list, np.ndarray)):
        return np.array(emb_val, dtype=np.float32)
    else:
        raise ValueError(f"Unknown embedding type: {type(emb_val)}")


class RouterDataset(Dataset):
    """Dataset for router training with difficulty information."""

    def __init__(
        self,
        parquet_path: str,
        embedding_dim: int = 512,
        embedding_prefix: str = "query_embedding_qwen3_instruct_router_v1",
        use_task_embedding: bool = True,
    ) -> None:
        """Initialize dataset from parquet file."""
        logger.info(f"Loading data from {parquet_path}...")
        df = pd.read_parquet(parquet_path)

        # Get unique LLMs and create mapping
        self.llm_names = sorted(df["llm"].unique().tolist())
        self.llm_to_idx = {name: i for i, name in enumerate(self.llm_names)}
        self.num_models = len(self.llm_names)
        logger.info(f"Found {self.num_models} LLMs: {self.llm_names}")

        # Get embedding column names
        query_emb_col = f"{embedding_prefix}_{embedding_dim}"
        task_emb_col = f"task_description_embedding_qwen3_instruct_router_v1_{embedding_dim}"

        if query_emb_col not in df.columns:
            raise ValueError(f"Embedding column {query_emb_col} not found")

        # Group by query to get unique queries
        logger.info("Grouping by query...")
        grouped = df.groupby(["task_id", "query"])

        # Extract embeddings, targets, and metadata
        query_embs = []
        task_embs = []
        targets = []
        n_correct_list = []  # Number of models correct (for difficulty)
        task_ids = []
        query_texts = []

        for (task_id, query), group in grouped:
            if len(group) != self.num_models:
                continue  # Skip incomplete groups

            # Get query embedding
            emb = parse_embedding(group[query_emb_col].iloc[0])
            query_embs.append(emb)

            # Get task embedding if available
            if use_task_embedding and task_emb_col in df.columns:
                task_emb = parse_embedding(group[task_emb_col].iloc[0])
                task_embs.append(task_emb)

            # Get effect for each LLM
            target = np.zeros(self.num_models)
            for _, row in group.iterrows():
                llm_idx = self.llm_to_idx[row["llm"]]
                target[llm_idx] = row["effect"]
            targets.append(target)

            # Count correct models (for difficulty)
            n_correct = (target > 0.5).sum()
            n_correct_list.append(n_correct)

            task_ids.append(task_id)
            query_texts.append(query)

        self.query_embs = torch.tensor(np.array(query_embs), dtype=torch.float32)
        self.targets = torch.tensor(np.array(targets), dtype=torch.float32)
        self.n_correct = np.array(n_correct_list)
        self.task_ids = task_ids
        self.query_texts = query_texts

        # Concatenate task embeddings if used
        if task_embs:
            task_embs_tensor = torch.tensor(np.array(task_embs), dtype=torch.float32)
            self.query_embs = torch.cat([self.query_embs, task_embs_tensor], dim=-1)
            self.embedding_dim = self.query_embs.shape[-1]
        else:
            self.embedding_dim = embedding_dim

        # Create model embeddings
        torch.manual_seed(42)
        self.model_embs = F.normalize(
            torch.randn(self.num_models, self.embedding_dim), dim=-1
        )

        # Assign difficulty bands
        self.difficulty_bands = self._assign_difficulty_bands()

        logger.info(f"Loaded {len(self)} queries with embedding dim {self.embedding_dim}")
        logger.info(f"Difficulty distribution: {self._get_difficulty_distribution()}")

    def _assign_difficulty_bands(self) -> list[str]:
        """Assign difficulty band to each query."""
        bands = []
        for n in self.n_correct:
            for band_name, (low, high) in DIFFICULTY_BANDS.items():
                if low <= n <= high:
                    bands.append(band_name)
                    break
            else:
                bands.append("unknown")
        return bands

    def _get_difficulty_distribution(self) -> dict[str, int]:
        """Get count of queries per difficulty band."""
        dist = defaultdict(int)
        for band in self.difficulty_bands:
            dist[band] += 1
        return dict(dist)

    def __len__(self) -> int:
        return len(self.query_embs)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.query_embs[idx], self.targets[idx]


# ==============================================================================
# Baseline Router Implementations
# ==============================================================================


class ResidualBlock(nn.Module):
    """Residual MLP block with LayerNorm."""

    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout),
        )
        self.norm = nn.LayerNorm(dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.norm(x + self.net(x))


class SEBlock(nn.Module):
    """Squeeze-and-Excitation block for channel recalibration."""

    def __init__(self, dim: int, reduction: int = 4) -> None:
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(dim, dim // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(dim // reduction, dim),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scale = self.fc(x)
        return x * scale


class TwoTowerRouter(nn.Module):
    """Two Tower architecture with separate query and model encoders."""

    def __init__(
        self,
        query_dim: int,
        model_dim: int,
        hidden_dim: int = 384,
        num_layers: int = 2,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        # Query tower
        self.query_proj = nn.Linear(query_dim, hidden_dim)
        self.query_blocks = nn.ModuleList([
            ResidualBlock(hidden_dim, hidden_dim * 2, dropout)
            for _ in range(num_layers)
        ])
        self.query_se = SEBlock(hidden_dim)
        self.query_norm = nn.LayerNorm(hidden_dim)

        # Model tower
        self.model_proj = nn.Linear(model_dim, hidden_dim)
        self.model_blocks = nn.ModuleList([
            ResidualBlock(hidden_dim, hidden_dim * 2, dropout)
            for _ in range(num_layers)
        ])
        self.model_se = SEBlock(hidden_dim)
        self.model_norm = nn.LayerNorm(hidden_dim)

        self.temperature = nn.Parameter(torch.ones(1))

    def encode_query(self, x: torch.Tensor) -> torch.Tensor:
        x = self.query_proj(x)
        for block in self.query_blocks:
            x = block(x)
        x = self.query_se(x)
        return self.query_norm(x)

    def encode_models(self, x: torch.Tensor) -> torch.Tensor:
        x = self.model_proj(x)
        for block in self.model_blocks:
            x = block(x)
        x = self.model_se(x)
        return self.model_norm(x)

    def forward(
        self, query_embs: torch.Tensor, model_embs: torch.Tensor
    ) -> torch.Tensor:
        query_repr = F.normalize(self.encode_query(query_embs), dim=-1)
        model_repr = F.normalize(self.encode_models(model_embs), dim=-1)
        scores = torch.matmul(query_repr, model_repr.T) / self.temperature.clamp(min=0.01)
        return scores

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


class MultiHeadMLPRouter(nn.Module):
    """Multi-Head MLP with shared encoder and per-model heads."""

    def __init__(
        self,
        query_dim: int,
        num_models: int,
        hidden_dim: int = 384,
        num_layers: int = 2,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        # Shared encoder
        self.encoder_proj = nn.Linear(query_dim, hidden_dim)
        self.encoder_blocks = nn.ModuleList([
            ResidualBlock(hidden_dim, hidden_dim * 2, dropout)
            for _ in range(num_layers)
        ])
        self.encoder_se = SEBlock(hidden_dim)
        self.encoder_norm = nn.LayerNorm(hidden_dim)

        # Per-model prediction heads
        self.heads = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim // 2),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim // 2, 1),
            )
            for _ in range(num_models)
        ])

    def forward(
        self, query_embs: torch.Tensor, model_embs: torch.Tensor = None
    ) -> torch.Tensor:
        # Encode query
        x = self.encoder_proj(query_embs)
        for block in self.encoder_blocks:
            x = block(x)
        x = self.encoder_se(x)
        x = self.encoder_norm(x)

        # Get predictions from each head
        scores = torch.cat([head(x) for head in self.heads], dim=-1)
        return scores

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


class SimplifiedGNNRouter(nn.Module):
    """Simplified GNN using PyTorch Geometric GAT layers."""

    def __init__(
        self,
        query_dim: int,
        model_dim: int,
        hidden_dim: int = 384,
        num_layers: int = 2,
        num_heads: int = 4,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        try:
            from torch_geometric.nn import GATConv
        except ImportError:
            raise ImportError("torch_geometric required for GNN router")

        # Input projections
        self.query_proj = nn.Linear(query_dim, hidden_dim)
        self.model_proj = nn.Linear(model_dim, hidden_dim)

        # GAT layers
        self.gat_layers = nn.ModuleList()
        for i in range(num_layers):
            in_dim = hidden_dim if i == 0 else hidden_dim * num_heads
            self.gat_layers.append(
                GATConv(in_dim, hidden_dim, heads=num_heads, dropout=dropout)
            )

        # Output projection
        final_dim = hidden_dim * num_heads
        self.output_proj = nn.Sequential(
            nn.Linear(final_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

        self.dropout = nn.Dropout(dropout)
        self.num_heads = num_heads
        self.hidden_dim = hidden_dim

    def _build_bipartite_graph(
        self, num_queries: int, num_models: int, device: torch.device
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build fully connected bipartite graph."""
        query_indices = torch.arange(num_queries, device=device)
        model_indices = torch.arange(num_models, device=device) + num_queries

        # Create edges: every query connects to every model
        src = query_indices.repeat_interleave(num_models)
        dst = model_indices.repeat(num_queries)

        edge_index = torch.stack([src, dst])
        return edge_index, query_indices, model_indices

    def forward(
        self, query_embs: torch.Tensor, model_embs: torch.Tensor
    ) -> torch.Tensor:
        batch_size = query_embs.shape[0]
        num_models = model_embs.shape[0]
        device = query_embs.device

        # Project inputs
        query_h = self.query_proj(query_embs)
        model_h = self.model_proj(model_embs)

        # Build graph
        edge_index, query_idx, model_idx = self._build_bipartite_graph(
            batch_size, num_models, device
        )

        # Combine node features
        x = torch.cat([query_h, model_h], dim=0)

        # Apply GAT layers
        for gat in self.gat_layers:
            x = self.dropout(F.elu(gat(x, edge_index)))

        # Extract query and model representations
        query_out = x[:batch_size]
        model_out = x[batch_size:]

        # Compute pairwise scores
        query_expanded = query_out.unsqueeze(1).expand(-1, num_models, -1)
        model_expanded = model_out.unsqueeze(0).expand(batch_size, -1, -1)
        combined = torch.cat([query_expanded, model_expanded], dim=-1)

        scores = self.output_proj(combined).squeeze(-1)
        return scores

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ==============================================================================
# Training
# ==============================================================================


def train_router(
    router: nn.Module,
    train_dataset: Dataset,
    val_dataset: Dataset,
    model_embs: torch.Tensor,
    num_epochs: int = 100,
    batch_size: int = 32,
    learning_rate: float = 1e-3,
    early_stopping_patience: int = 20,
    device: str = "cuda",
) -> tuple[dict, float]:
    """Train a router and return training history."""
    router = router.to(device)
    model_embs = model_embs.to(device)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)

    optimizer = torch.optim.AdamW(router.parameters(), lr=learning_rate, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)

    history = {
        "train_loss": [],
        "val_loss": [],
        "val_frac_oracle": [],
    }

    best_val_frac_oracle = 0
    best_epoch = 0
    best_state = None
    patience_counter = 0

    start_time = time.time()

    for epoch in range(num_epochs):
        # Training
        router.train()
        train_losses = []
        for query_batch, target_batch in train_loader:
            query_batch = query_batch.to(device)
            target_batch = target_batch.to(device)

            optimizer.zero_grad()
            scores = router(query_batch, model_embs)

            # Hybrid loss: contrastive + KL
            log_probs = F.log_softmax(scores, dim=-1)
            target_probs = F.softmax(target_batch / 0.5, dim=-1)
            kl_loss = F.kl_div(log_probs, target_probs, reduction="batchmean")

            # Contrastive component
            best_idx = target_batch.argmax(dim=-1)
            contrastive_loss = F.cross_entropy(scores, best_idx)

            loss = 0.5 * kl_loss + 0.5 * contrastive_loss
            loss.backward()

            torch.nn.utils.clip_grad_norm_(router.parameters(), 1.0)
            optimizer.step()
            train_losses.append(loss.item())

        scheduler.step()

        # Validation
        router.eval()
        val_losses = []
        all_preds = []
        all_targets = []

        with torch.no_grad():
            for query_batch, target_batch in val_loader:
                query_batch = query_batch.to(device)
                target_batch = target_batch.to(device)

                scores = router(query_batch, model_embs)

                log_probs = F.log_softmax(scores, dim=-1)
                target_probs = F.softmax(target_batch / 0.5, dim=-1)
                kl_loss = F.kl_div(log_probs, target_probs, reduction="batchmean")
                best_idx = target_batch.argmax(dim=-1)
                contrastive_loss = F.cross_entropy(scores, best_idx)
                loss = 0.5 * kl_loss + 0.5 * contrastive_loss

                val_losses.append(loss.item())
                all_preds.append(scores.cpu())
                all_targets.append(target_batch.cpu())

        all_preds = torch.cat(all_preds)
        all_targets = torch.cat(all_targets)

        # Compute frac_oracle
        pred_idx = all_preds.argmax(dim=-1)
        oracle_idx = all_targets.argmax(dim=-1)
        pred_rewards = all_targets.gather(1, pred_idx.unsqueeze(1)).squeeze()
        oracle_rewards = all_targets.gather(1, oracle_idx.unsqueeze(1)).squeeze()
        frac_oracle = (pred_rewards / oracle_rewards.clamp(min=1e-6)).mean().item()

        history["train_loss"].append(np.mean(train_losses))
        history["val_loss"].append(np.mean(val_losses))
        history["val_frac_oracle"].append(frac_oracle)

        if frac_oracle > best_val_frac_oracle:
            best_val_frac_oracle = frac_oracle
            best_epoch = epoch + 1
            best_state = {k: v.cpu().clone() for k, v in router.state_dict().items()}
            patience_counter = 0
        else:
            patience_counter += 1

        if (epoch + 1) % 10 == 0:
            logger.info(
                f"  Epoch {epoch+1}/{num_epochs} - "
                f"Loss: {history['train_loss'][-1]:.4f}, "
                f"Val Loss: {history['val_loss'][-1]:.4f}, "
                f"FracOracle: {frac_oracle:.4f}"
            )

        if patience_counter >= early_stopping_patience:
            logger.info(f"  Early stopping at epoch {epoch+1}")
            break

    training_time = time.time() - start_time

    # Restore best state
    if best_state is not None:
        router.load_state_dict(best_state)

    history["best_epoch"] = best_epoch
    return history, training_time


# ==============================================================================
# Evaluation Metrics
# ==============================================================================


def compute_difficulty_metrics(
    preds: torch.Tensor,
    targets: torch.Tensor,
    n_correct: np.ndarray,
    difficulty_bands: list[str],
) -> list[DifficultyMetrics]:
    """Compute metrics stratified by difficulty band."""
    results = []

    pred_idx = preds.argmax(dim=-1)
    oracle_idx = targets.argmax(dim=-1)

    for band_name in DIFFICULTY_BANDS.keys():
        mask = np.array([b == band_name for b in difficulty_bands])
        n_band = mask.sum()

        if n_band == 0:
            continue

        band_preds = preds[mask]
        band_targets = targets[mask]
        band_pred_idx = pred_idx[mask]
        band_oracle_idx = oracle_idx[mask]

        # Top-1 accuracy
        top1_acc = (band_pred_idx == band_oracle_idx).float().mean().item()

        # Frac of oracle
        pred_rewards = band_targets.gather(1, band_pred_idx.unsqueeze(1)).squeeze()
        oracle_rewards = band_targets.gather(1, band_oracle_idx.unsqueeze(1)).squeeze()
        frac_oracle = (pred_rewards / oracle_rewards.clamp(min=1e-6)).mean().item()

        # Regret
        regret = (oracle_rewards - pred_rewards).mean().item()

        # Average effect
        avg_effect = band_targets.mean().item()

        results.append(DifficultyMetrics(
            band_name=band_name,
            n_queries=int(n_band),
            pct_of_total=n_band / len(difficulty_bands),
            top1_acc=top1_acc,
            frac_of_oracle=frac_oracle,
            avg_regret=regret,
            avg_effect=avg_effect,
        ))

    return results


def compute_calibration_metrics(
    preds: torch.Tensor,
    targets: torch.Tensor,
) -> CalibrationMetrics:
    """Compute uncertainty calibration metrics."""
    probs = F.softmax(preds, dim=-1)
    max_prob, pred_idx = probs.max(dim=-1)
    oracle_idx = targets.argmax(dim=-1)

    # Is prediction correct?
    correct = (pred_idx == oracle_idx).float()
    errors = 1 - correct

    # Entropy
    entropy = -(probs * probs.clamp(min=1e-10).log()).sum(dim=-1)
    max_entropy = np.log(preds.shape[1])
    normalized_entropy = entropy / max_entropy

    # Frac oracle per query
    pred_rewards = targets.gather(1, pred_idx.unsqueeze(1)).squeeze()
    oracle_rewards = targets.gather(1, oracle_idx.unsqueeze(1)).squeeze()
    frac_oracle = pred_rewards / oracle_rewards.clamp(min=1e-6)

    # ECE (Expected Calibration Error)
    n_bins = 10
    bin_boundaries = torch.linspace(0, 1, n_bins + 1)
    ece = 0
    for i in range(n_bins):
        mask = (max_prob >= bin_boundaries[i]) & (max_prob < bin_boundaries[i + 1])
        if mask.sum() > 0:
            bin_conf = max_prob[mask].mean()
            bin_acc = correct[mask].mean()
            ece += mask.sum() * torch.abs(bin_conf - bin_acc)
    ece = (ece / len(max_prob)).item()

    # Correlations
    entropy_np = normalized_entropy.numpy()
    error_np = errors.numpy()
    frac_oracle_np = frac_oracle.numpy()

    entropy_error_corr = np.corrcoef(entropy_np, error_np)[0, 1]
    entropy_quality_corr = np.corrcoef(entropy_np, frac_oracle_np)[0, 1]

    # Logit margin
    sorted_logits = preds.sort(dim=-1, descending=True).values
    margin = (sorted_logits[:, 0] - sorted_logits[:, 1]).numpy()
    margin_error_corr = np.corrcoef(margin, error_np)[0, 1]

    return CalibrationMetrics(
        ece=ece,
        entropy_error_corr=float(np.nan_to_num(entropy_error_corr)),
        entropy_quality_corr=float(np.nan_to_num(entropy_quality_corr)),
        logit_margin_error_corr=float(np.nan_to_num(margin_error_corr)),
        avg_entropy=normalized_entropy.mean().item(),
        avg_max_prob=max_prob.mean().item(),
    )


def compute_error_analysis(
    preds: torch.Tensor,
    targets: torch.Tensor,
    difficulty_bands: list[str],
    llm_names: list[str],
) -> ErrorAnalysis:
    """Compute detailed error analysis."""
    pred_idx = preds.argmax(dim=-1)
    oracle_idx = targets.argmax(dim=-1)

    # Is prediction correct?
    errors = pred_idx != oracle_idx

    # Error regret
    pred_rewards = targets.gather(1, pred_idx.unsqueeze(1)).squeeze()
    oracle_rewards = targets.gather(1, oracle_idx.unsqueeze(1)).squeeze()
    regret = oracle_rewards - pred_rewards

    n_errors = errors.sum().item()
    error_rate = n_errors / len(errors)
    avg_error_regret = regret[errors].mean().item() if n_errors > 0 else 0

    # Errors by band
    errors_by_band = {}
    error_rate_by_band = {}
    for band_name in DIFFICULTY_BANDS.keys():
        mask = np.array([b == band_name for b in difficulty_bands])
        n_band = mask.sum()
        if n_band > 0:
            band_errors = errors.numpy()[mask].sum()
            errors_by_band[band_name] = int(band_errors)
            error_rate_by_band[band_name] = band_errors / n_band
        else:
            errors_by_band[band_name] = 0
            error_rate_by_band[band_name] = 0.0

    # Worst errors
    error_regrets = regret[errors].numpy()
    worst_indices = np.argsort(error_regrets)[-10:][::-1]
    worst_error_regrets = error_regrets[worst_indices].tolist()

    # Confusion matrix
    confusion = defaultdict(lambda: defaultdict(int))
    for i in range(len(pred_idx)):
        if errors[i]:
            pred_llm = llm_names[pred_idx[i]]
            actual_llm = llm_names[oracle_idx[i]]
            confusion[pred_llm][actual_llm] += 1

    # Most confused pairs
    pairs = []
    for pred_llm, actuals in confusion.items():
        for actual_llm, count in actuals.items():
            pairs.append((pred_llm, actual_llm, count))
    pairs.sort(key=lambda x: x[2], reverse=True)
    most_confused = pairs[:5]

    return ErrorAnalysis(
        n_errors=int(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,
        confusion_matrix=dict(confusion),
        most_confused_pairs=most_confused,
    )


def compute_ranking_metrics(
    preds: torch.Tensor,
    targets: torch.Tensor,
) -> RankingMetrics:
    """Compute ranking quality metrics."""
    n_queries = preds.shape[0]

    spearman_corrs = []
    kendall_taus = []

    for i in range(n_queries):
        pred_ranks = preds[i].numpy()
        true_ranks = targets[i].numpy()

        spearman = stats.spearmanr(pred_ranks, true_ranks).correlation
        kendall = stats.kendalltau(pred_ranks, true_ranks).correlation

        if not np.isnan(spearman):
            spearman_corrs.append(spearman)
        if not np.isnan(kendall):
            kendall_taus.append(kendall)

    avg_spearman = np.mean(spearman_corrs) if spearman_corrs else 0
    avg_kendall = np.mean(kendall_taus) if kendall_taus else 0
    pct_positive = np.mean([s > 0 for s in spearman_corrs]) if spearman_corrs else 0

    # NDCG
    def dcg_at_k(scores, k):
        return np.sum(scores[:k] / np.log2(np.arange(2, k + 2)))

    ndcg_1_list = []
    ndcg_3_list = []

    for i in range(n_queries):
        pred_order = preds[i].argsort(descending=True)
        true_relevance = targets[i].numpy()
        sorted_relevance = true_relevance[pred_order]

        ideal_relevance = np.sort(true_relevance)[::-1]

        dcg1 = dcg_at_k(sorted_relevance, 1)
        idcg1 = dcg_at_k(ideal_relevance, 1)
        ndcg_1_list.append(dcg1 / idcg1 if idcg1 > 0 else 0)

        dcg3 = dcg_at_k(sorted_relevance, 3)
        idcg3 = dcg_at_k(ideal_relevance, 3)
        ndcg_3_list.append(dcg3 / idcg3 if idcg3 > 0 else 0)

    return RankingMetrics(
        avg_spearman_corr=avg_spearman,
        avg_kendall_tau=avg_kendall,
        pct_positive_corr=pct_positive,
        ndcg_at_1=np.mean(ndcg_1_list),
        ndcg_at_3=np.mean(ndcg_3_list),
    )


def evaluate_router(
    router: nn.Module,
    test_dataset: Dataset,
    model_embs: torch.Tensor,
    llm_names: list[str],
    difficulty_bands: list[str],
    n_correct: np.ndarray,
    device: str = "cuda",
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run evaluation and return predictions and targets."""
    router = router.to(device)
    router.eval()
    model_embs = model_embs.to(device)

    test_loader = DataLoader(test_dataset, batch_size=64)

    all_preds = []
    all_targets = []

    with torch.no_grad():
        for query_batch, target_batch in test_loader:
            query_batch = query_batch.to(device)
            scores = router(query_batch, model_embs)
            all_preds.append(scores.cpu())
            all_targets.append(target_batch)

    preds = torch.cat(all_preds)
    targets = torch.cat(all_targets)

    return preds, targets


# ==============================================================================
# Main Entry Point
# ==============================================================================


def run_comprehensive_evaluation(
    data_path: str,
    output_dir: str,
    embedding_dims: list[int] = [256, 512, 1024],
    num_epochs: int = 100,
    batch_size: int = 32,
    device: str = "cuda",
) -> dict:
    """Run comprehensive evaluation across all baselines and embedding dims."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    all_results = []

    for embed_dim in embedding_dims:
        logger.info(f"\n{'='*80}")
        logger.info(f"EMBEDDING DIMENSION: {embed_dim}")
        logger.info(f"{'='*80}")

        # Load dataset
        dataset = RouterDataset(data_path, embedding_dim=embed_dim)

        # Split data
        n = len(dataset)
        train_size = int(0.7 * n)
        val_size = int(0.15 * n)
        test_size = n - train_size - val_size

        indices = torch.randperm(n).tolist()
        train_indices = indices[:train_size]
        val_indices = indices[train_size:train_size + val_size]
        test_indices = indices[train_size + val_size:]

        train_dataset = Subset(dataset, train_indices)
        val_dataset = Subset(dataset, val_indices)
        test_dataset = Subset(dataset, test_indices)

        # Get metadata for test set
        test_difficulty_bands = [dataset.difficulty_bands[i] for i in test_indices]
        test_n_correct = dataset.n_correct[test_indices]

        logger.info(f"Data split: train={len(train_dataset)}, val={len(val_dataset)}, test={len(test_dataset)}")

        # Create routers
        routers = [
            ("Two Tower", TwoTowerRouter(
                query_dim=dataset.embedding_dim,
                model_dim=dataset.embedding_dim,
                hidden_dim=384,
                num_layers=2,
                dropout=0.1,
            )),
            ("Multi-Head MLP", MultiHeadMLPRouter(
                query_dim=dataset.embedding_dim,
                num_models=dataset.num_models,
                hidden_dim=384,
                num_layers=2,
                dropout=0.1,
            )),
            ("GNN (Simplified)", SimplifiedGNNRouter(
                query_dim=dataset.embedding_dim,
                model_dim=dataset.embedding_dim,
                hidden_dim=384,
                num_layers=2,
                num_heads=4,
                dropout=0.1,
            )),
        ]

        for name, router in routers:
            logger.info(f"\n{'-'*60}")
            logger.info(f"Training {name} (embed_dim={embed_dim}, params={router.count_parameters():,})")
            logger.info(f"{'-'*60}")

            # Train
            history, training_time = train_router(
                router=router,
                train_dataset=train_dataset,
                val_dataset=val_dataset,
                model_embs=dataset.model_embs,
                num_epochs=num_epochs,
                batch_size=batch_size,
                device=device,
            )

            # Save checkpoint
            checkpoint_path = output_dir / f"{name.replace(' ', '_').lower()}_dim{embed_dim}.pt"
            torch.save({
                "state_dict": router.state_dict(),
                "config": {
                    "name": name,
                    "embedding_dim": embed_dim,
                    "params": router.count_parameters(),
                },
                "history": history,
            }, checkpoint_path)
            logger.info(f"  Saved checkpoint to {checkpoint_path}")

            # Evaluate
            preds, targets = evaluate_router(
                router=router,
                test_dataset=test_dataset,
                model_embs=dataset.model_embs,
                llm_names=dataset.llm_names,
                difficulty_bands=test_difficulty_bands,
                n_correct=test_n_correct,
                device=device,
            )

            # Core metrics
            pred_idx = preds.argmax(dim=-1)
            oracle_idx = targets.argmax(dim=-1)
            top1_acc = (pred_idx == oracle_idx).float().mean().item()

            pred_rewards = targets.gather(1, pred_idx.unsqueeze(1)).squeeze()
            oracle_rewards = targets.gather(1, oracle_idx.unsqueeze(1)).squeeze()
            frac_oracle = (pred_rewards / oracle_rewards.clamp(min=1e-6)).mean().item()
            regret = (oracle_rewards - pred_rewards).mean().item()

            # Difficulty-aware metrics
            difficulty_metrics = compute_difficulty_metrics(
                preds, targets, test_n_correct, test_difficulty_bands
            )

            # Calibration
            calibration = compute_calibration_metrics(preds, targets)

            # Error analysis
            errors = compute_error_analysis(
                preds, targets, test_difficulty_bands, dataset.llm_names
            )

            # Ranking metrics
            ranking = compute_ranking_metrics(preds, targets)

            # Compile result
            result = BaselineResult(
                name=name,
                embedding_dim=embed_dim,
                params=router.count_parameters(),
                training_time=training_time,
                best_epoch=history["best_epoch"],
                test_frac_oracle=frac_oracle,
                test_top1_acc=top1_acc,
                test_regret=regret,
                train_loss_history=history["train_loss"],
                val_loss_history=history["val_loss"],
                val_frac_oracle_history=history["val_frac_oracle"],
                difficulty_metrics=difficulty_metrics,
                calibration=calibration,
                errors=errors,
                ranking=ranking,
            )

            all_results.append(result)

            # Log summary
            logger.info(f"\n  === TEST RESULTS ===")
            logger.info(f"  FracOracle: {frac_oracle:.4f}")
            logger.info(f"  Top1 Acc: {top1_acc:.4f}")
            logger.info(f"  Regret: {regret:.4f}")
            logger.info(f"  Spearman: {ranking.avg_spearman_corr:.4f}")
            logger.info(f"  NDCG@1: {ranking.ndcg_at_1:.4f}")
            logger.info(f"  ECE: {calibration.ece:.4f}")
            logger.info(f"  Entropy-Error Corr: {calibration.entropy_error_corr:.4f}")

            logger.info(f"\n  === DIFFICULTY BREAKDOWN ===")
            for dm in difficulty_metrics:
                logger.info(
                    f"  {dm.band_name:12s}: n={dm.n_queries:5d}, "
                    f"FracOracle={dm.frac_of_oracle:.4f}, "
                    f"Top1={dm.top1_acc:.4f}, "
                    f"Regret={dm.avg_regret:.4f}"
                )

    # Save all results
    results_path = output_dir / "comprehensive_results.json"
    with open(results_path, "w") as f:
        json.dump([r.to_dict() for r in all_results], f, indent=2)
    logger.info(f"\nSaved comprehensive results to {results_path}")

    # Print summary table
    print_summary_table(all_results)

    return all_results


def print_summary_table(results: list[BaselineResult]) -> None:
    """Print summary comparison table."""
    logger.info("\n" + "=" * 100)
    logger.info("COMPREHENSIVE BASELINE COMPARISON")
    logger.info("=" * 100)

    # Group by embedding dim
    by_dim = defaultdict(list)
    for r in results:
        by_dim[r.embedding_dim].append(r)

    for dim in sorted(by_dim.keys()):
        logger.info(f"\n--- Embedding Dimension: {dim} ---")
        logger.info(
            f"{'Model':<20} {'Params':>10} {'FracOracle':>12} {'Top1':>8} "
            f"{'Spearman':>10} {'NDCG@1':>8} {'ECE':>8} {'Time':>8}"
        )
        logger.info("-" * 100)

        for r in by_dim[dim]:
            logger.info(
                f"{r.name:<20} {r.params:>10,} {r.test_frac_oracle:>12.4f} "
                f"{r.test_top1_acc:>8.4f} {r.ranking.avg_spearman_corr:>10.4f} "
                f"{r.ranking.ndcg_at_1:>8.4f} {r.calibration.ece:>8.4f} "
                f"{r.training_time:>7.1f}s"
            )

    # Best overall
    logger.info("\n" + "=" * 100)
    logger.info("BEST MODELS")
    logger.info("=" * 100)

    best_frac = max(results, key=lambda r: r.test_frac_oracle)
    best_spearman = max(results, key=lambda r: r.ranking.avg_spearman_corr)
    best_ece = min(results, key=lambda r: r.calibration.ece)

    logger.info(f"Best FracOracle: {best_frac.name} (dim={best_frac.embedding_dim}) = {best_frac.test_frac_oracle:.4f}")
    logger.info(f"Best Spearman: {best_spearman.name} (dim={best_spearman.embedding_dim}) = {best_spearman.ranking.avg_spearman_corr:.4f}")
    logger.info(f"Best Calibration: {best_ece.name} (dim={best_ece.embedding_dim}) = ECE {best_ece.calibration.ece:.4f}")


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Comprehensive baseline evaluation")
    parser.add_argument(
        "--data_path",
        type=str,
        default="embergraphrouter/data/router_data_8llm_qwen3.parquet",
        help="Path to router data parquet",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="results/baselines_comprehensive",
        help="Output directory",
    )
    parser.add_argument(
        "--embedding_dims",
        type=int,
        nargs="+",
        default=[256, 512, 1024],
        help="Embedding dimensions to test",
    )
    parser.add_argument(
        "--num_epochs",
        type=int,
        default=100,
        help="Maximum training epochs",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=32,
        help="Batch size",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device (cuda or cpu)",
    )

    args = parser.parse_args()

    logger.info(f"Device: {args.device}")
    if args.device == "cuda":
        logger.info(f"GPU: {torch.cuda.get_device_name()}")

    run_comprehensive_evaluation(
        data_path=args.data_path,
        output_dir=args.output_dir,
        embedding_dims=args.embedding_dims,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        device=args.device,
    )


if __name__ == "__main__":
    main()
