"""EB1-GNN Tiered Inference Router.

This module implements intelligent, difficulty-aware routing that uses GNN predictions
to select the optimal inference strategy:

Tier 1 - TRIVIAL (Easy):
    All models predicted to succeed with high confidence. Route to the cheapest/fastest
    model (e.g., Gemini 2.5 Flash) to minimize cost without sacrificing quality.

Tier 2 - DISCRIMINATIVE (Medium):
    Models differ in predicted performance but GNN is confident in its ranking.
    Route to GNN's top choice or a strong model (e.g., Gemini 3 Pro).

Tier 3 - UNCERTAIN (Hard):
    High disagreement and/or low confidence. Use full EB1 ensemble for maximum
    quality on challenging queries, or EB1-Pro for critical applications.

The threshold strategy is based on empirical analysis from eval_difficulty_aware.py:
- mean_logit signal captures P(any LLM correct) better than softmax entropy
- logit_margin indicates decisiveness of the routing decision
- Combined with perf_range prediction for cost-benefit analysis

Example:
    from embergraphrouter.inference import TieredRouter

    router = TieredRouter.from_registry("eb1-gnn-preview-12-19")

    # Automatic tier selection
    result = router.route("What is 2 + 2?")  # -> TRIVIAL tier
    result = router.route("Prove Fermat's Last Theorem")  # -> UNCERTAIN tier

    # With custom thresholds
    router = TieredRouter.from_registry(
        "eb1-gnn-preview-12-19",
        thresholds=TieredThresholds(
            trivial_mean_logit=0.6,  # More conservative trivial detection
            uncertain_max_prob=0.25,  # Lower uncertainty trigger
        ),
    )
"""

from __future__ import annotations

import logging
import random
import time
from dataclasses import dataclass, field
from enum import Enum, auto
from pathlib import Path
from typing import TYPE_CHECKING, Mapping, Sequence

import numpy as np
import torch
import torch.nn.functional as F

if TYPE_CHECKING:
    from ember.models.schemas import ChatResponse


logger = logging.getLogger(__name__)


class InferenceTier(Enum):
    """Inference tier based on query difficulty."""

    TRIVIAL = auto()  # Easy: all models correct, use cheapest
    DISCRIMINATIVE = auto()  # Medium: GNN routing matters
    UNCERTAIN = auto()  # Hard: ensemble for maximum quality
    OUT_OF_DISTRIBUTION = auto()  # OOD: query outside training distribution


def _load_default_thresholds() -> dict:
    """Load default thresholds from central YAML config.

    Returns:
        Dictionary of threshold values, or empty dict if config not found.
    """
    import yaml

    default_config_path = Path(__file__).parent.parent / "registry" / "default_thresholds.yaml"
    if not default_config_path.exists():
        return {}

    try:
        with open(default_config_path) as f:
            data = yaml.safe_load(f)
        tiered_config = data.get("tiered_routing", {})
        thresholds = tiered_config.get("thresholds", {})
        ood_config = tiered_config.get("ood_detection", {})
        return {
            "trivial_mean_logit": thresholds.get("trivial_mean_logit"),
            "uncertain_mean_logit": thresholds.get("uncertain_mean_logit"),
            "uncertain_max_prob": thresholds.get("uncertain_max_prob"),
            "uncertain_logit_margin": thresholds.get("uncertain_logit_margin"),
            "min_ensemble_value": thresholds.get("min_ensemble_value"),
            "ood_max_confidence": ood_config.get("max_confidence"),
            "ood_max_logit_std": ood_config.get("max_logit_std"),
            "ood_detection_enabled": ood_config.get("enabled"),
            "ood_ensemble_probability": ood_config.get("ensemble_probability"),
            "ood_contrastive_threshold": ood_config.get("contrastive_threshold"),
        }
    except Exception:
        return {}


# Cache default thresholds at module load time
_DEFAULT_THRESHOLDS = _load_default_thresholds()


def _get_default(key: str, fallback):
    """Get a default threshold value from config, with fallback."""
    value = _DEFAULT_THRESHOLDS.get(key)
    return value if value is not None else fallback


@dataclass(frozen=True, slots=True)
class TieredThresholds:
    """Thresholds for tier classification.

    These thresholds are calibrated based on empirical analysis from
    eval_difficulty_aware.py and offline calibration. The key insight is that
    mean_logit (average raw logit value) captures difficulty better than
    softmax entropy or max_prob, which compress differences into near-uniform
    probabilities with 7+ LLMs.

    IMPORTANT: Thresholds are model-specific and must be calibrated for each
    checkpoint. Use TieredThresholds.from_registry() to load calibrated values,
    or run embergraphrouter/scripts/run_threshold_calibration.py for a new model.

    Default values are loaded from registry/default_thresholds.yaml. This ensures
    a single source of truth for threshold values across the codebase.

    Calibration showed:
    - mean_logit correlation with n_correct: 0.606 (strong)
    - max_prob is nearly constant (~0.17-0.19) regardless of difficulty
    - TRIVIAL detection uses mean_logit only (max_prob check removed)

    Attributes:
        trivial_mean_logit: Queries with mean_logit above this are trivial.
        uncertain_mean_logit: Queries with mean_logit below this are uncertain.
        uncertain_max_prob: Queries with max_prob below this are uncertain.
            Secondary signal for cases where logits are unusually compressed.
        uncertain_logit_margin: Low margin between top choices indicates uncertainty.
            If top1_logit - top2_logit < this, the model can't distinguish.
        min_ensemble_value: Minimum perf_range to justify ensemble cost.
            If all models perform similarly, ensemble adds cost without benefit.
        ood_max_confidence: Queries with max_prob below this AND collapsed logits
            are likely out-of-distribution. Default 0.16 based on SWE-bench analysis.
        ood_max_logit_std: Queries with logit_std below this have collapsed
            predictions indicating OOD. Default 0.08 based on analysis.
        ood_detection_enabled: Whether to enable OOD detection. When disabled,
            OOD queries are handled via the normal tier classification.
        ood_ensemble_probability: Probability of using ensemble instead of top-tier
            for OOD queries (exploration vs exploitation). Default 0.0 (always top-tier).
            Set to 0.1-0.3 for exploration, or 1.0 to always use ensemble on OOD.
        ood_contrastive_threshold: Contrastive alignment threshold for OOD detection.
            Queries with max similarity to any LLM below this are OOD. A value of 0.0
            disables contrastive detection (falls back to logit stats).
    """

    trivial_mean_logit: float = _get_default("trivial_mean_logit", 0.5)
    uncertain_mean_logit: float = _get_default("uncertain_mean_logit", -0.5)
    uncertain_max_prob: float = _get_default("uncertain_max_prob", 0.15)
    uncertain_logit_margin: float = _get_default("uncertain_logit_margin", 0.1)
    min_ensemble_value: float = _get_default("min_ensemble_value", 0.15)
    ood_max_confidence: float = _get_default("ood_max_confidence", 0.16)
    ood_max_logit_std: float = _get_default("ood_max_logit_std", 0.08)
    ood_detection_enabled: bool = _get_default("ood_detection_enabled", True)
    ood_ensemble_probability: float = _get_default("ood_ensemble_probability", 0.0)
    ood_contrastive_threshold: float = _get_default("ood_contrastive_threshold", 0.0)

    @classmethod
    def from_registry(
        cls,
        registry_path: str | Path,
        *,
        fallback_to_defaults: bool = True,
    ) -> "TieredThresholds":
        """Load calibrated thresholds from registry's thresholds.yaml.

        Each GNN checkpoint produces a specific logit distribution. Thresholds
        must be calibrated for that distribution to work correctly. This method
        loads thresholds that were computed by run_threshold_calibration.py and
        stored alongside the model weights.

        Args:
            registry_path: Path to registry directory containing thresholds.yaml.
            fallback_to_defaults: If True and thresholds.yaml not found, return
                default thresholds with a warning. If False, raise FileNotFoundError.

        Returns:
            TieredThresholds with calibrated values.

        Raises:
            FileNotFoundError: If thresholds.yaml not found and
                fallback_to_defaults is False.
        """
        import yaml

        registry_path = Path(registry_path)
        thresholds_path = registry_path / "thresholds.yaml"

        if not thresholds_path.exists():
            if fallback_to_defaults:
                logger.warning(
                    "No thresholds.yaml found at %s. Using defaults. "
                    "Run run_threshold_calibration.py to generate calibrated "
                    "thresholds for this model checkpoint.",
                    thresholds_path,
                )
                return cls()
            raise FileNotFoundError(
                f"Calibrated thresholds not found: {thresholds_path}. "
                f"Run embergraphrouter/scripts/run_threshold_calibration.py "
                f"to generate thresholds for this model."
            )

        with open(thresholds_path) as f:
            data = yaml.safe_load(f)

        tiered_config = data.get("tiered_routing", {})
        thresholds_data = tiered_config.get("thresholds", {})

        logger.info(
            "Loaded calibrated thresholds from %s: trivial=%.3f, uncertain=%.3f, margin=%.4f",
            thresholds_path,
            thresholds_data.get("trivial_mean_logit", 0.5),
            thresholds_data.get("uncertain_mean_logit", -0.5),
            thresholds_data.get("uncertain_logit_margin", 0.1),
        )

        # Load OOD detection config
        ood_config = tiered_config.get("ood_detection", {})

        contrastive_threshold = ood_config.get("contrastive_threshold", 0.0)
        if contrastive_threshold > 0:
            logger.info(
                "Loaded contrastive OOD threshold: %.3f (primary OOD detection)",
                contrastive_threshold,
            )

        return cls(
            trivial_mean_logit=thresholds_data.get("trivial_mean_logit", 0.5),
            uncertain_mean_logit=thresholds_data.get("uncertain_mean_logit", -0.5),
            uncertain_max_prob=thresholds_data.get("uncertain_max_prob", 0.15),
            uncertain_logit_margin=thresholds_data.get("uncertain_logit_margin", 0.1),
            min_ensemble_value=thresholds_data.get("min_ensemble_value", 0.15),
            ood_max_confidence=ood_config.get("max_confidence", 0.16),
            ood_max_logit_std=ood_config.get("max_logit_std", 0.08),
            ood_detection_enabled=ood_config.get("enabled", True),
            ood_ensemble_probability=ood_config.get("ensemble_probability", 0.0),
            ood_contrastive_threshold=contrastive_threshold,
        )


@dataclass(frozen=True, slots=True)
class TierConfig:
    """Configuration for a specific inference tier.

    Attributes:
        tier: The inference tier this config applies to.
        model: Model identifier to use for this tier.
        model_params: Additional parameters for model invocation.
        use_ensemble: Whether to use EB1 ensemble instead of single model.
        ensemble_plan: If use_ensemble, which EB1 plan to use.
    """

    tier: InferenceTier
    model: str
    model_params: Mapping[str, object] = field(default_factory=dict)
    use_ensemble: bool = False
    ensemble_plan: str | None = None


# Default tier configurations matching production model availability
DEFAULT_TIER_CONFIGS = {
    InferenceTier.TRIVIAL: TierConfig(
        tier=InferenceTier.TRIVIAL,
        model="gemini-2.5-flash",  # Cheapest, still capable
        model_params={},
        use_ensemble=False,
    ),
    InferenceTier.DISCRIMINATIVE: TierConfig(
        tier=InferenceTier.DISCRIMINATIVE,
        model="gemini-3.1-pro-preview",  # Strong, cost-effective
        model_params={},
        use_ensemble=False,
    ),
    InferenceTier.UNCERTAIN: TierConfig(
        tier=InferenceTier.UNCERTAIN,
        model="",  # Ignored when use_ensemble=True
        model_params={},
        use_ensemble=True,
        ensemble_plan="DEFAULT_PLAN",  # Full EB1 ensemble
    ),
    InferenceTier.OUT_OF_DISTRIBUTION: TierConfig(
        tier=InferenceTier.OUT_OF_DISTRIBUTION,
        model="claude-opus-4-6",  # Top-tier for safety on OOD
        model_params={},
        use_ensemble=False,
    ),
}


@dataclass
class DifficultySignals:
    """Signals extracted from GNN prediction for tier classification.

    These signals are computed from the raw model outputs before and after
    softmax to capture different aspects of difficulty:

    - mean_logit: Average raw logit, correlates with P(any model correct)
    - max_prob: Confidence in top choice after softmax
    - logit_margin: How decisive the routing decision is
    - entropy: Information-theoretic uncertainty (less useful due to compression)
    - max_contrastive_similarity: Max similarity to any LLM in contrastive space
      (principled OOD detection - low similarity indicates out-of-distribution)
    """

    mean_logit: float
    max_prob: float
    logit_margin: float
    logit_std: float
    entropy: float
    top1_model: str
    top1_idx: int
    top2_model: str
    top2_idx: int
    all_probs: dict[str, float]
    all_logits: dict[str, float]
    # Contrastive alignment signal for principled OOD detection
    # None if model doesn't have contrastive encoder
    max_contrastive_similarity: float | None = None

    def classify_tier(self, thresholds: TieredThresholds) -> InferenceTier:
        """Classify query into inference tier based on signals.

        The classification logic implements a decision tree:
        0. Check for OOD: contrastive alignment (primary) or collapsed predictions (fallback)
        1. Check for TRIVIAL: high mean_logit (all models likely correct)
        2. Check for UNCERTAIN: low mean_logit OR low margin (need ensemble)
        3. Default to DISCRIMINATIVE: GNN routing is valuable

        Note: TRIVIAL detection uses only mean_logit because max_prob is nearly
        constant (~0.17-0.19) with 7 LLMs due to softmax compression. Calibration
        showed mean_logit has 0.606 correlation with query difficulty.

        OOD detection (updated 2025-01-25):
        - PRIMARY: Contrastive alignment - queries with low max similarity to any LLM
          are outside the training distribution. This is principled because the
          contrastive encoder is trained via InfoNCE to align ID queries with LLMs.
        - FALLBACK: Logit statistics - for models without contrastive encoder,
          near-uniform predictions + collapsed logits indicate OOD.

        Args:
            thresholds: Threshold configuration for tier boundaries.

        Returns:
            The appropriate inference tier.
        """
        # Tier 0: OUT_OF_DISTRIBUTION - check first as OOD may have arbitrary logits
        if thresholds.ood_detection_enabled:
            # PRIMARY: Contrastive alignment (principled, theoretically grounded)
            # Low max similarity to any LLM indicates the query was never trained
            # to align with any model - a direct measure of OOD status
            if (
                self.max_contrastive_similarity is not None
                and thresholds.ood_contrastive_threshold > 0
            ):
                if self.max_contrastive_similarity < thresholds.ood_contrastive_threshold:
                    return InferenceTier.OUT_OF_DISTRIBUTION
            else:
                # FALLBACK: Logit statistics (for models without contrastive encoder)
                # Near-uniform predictions + collapsed logits = no signal about query
                is_near_uniform = self.max_prob < thresholds.ood_max_confidence
                is_collapsed = self.logit_std < thresholds.ood_max_logit_std
                if is_near_uniform and is_collapsed:
                    return InferenceTier.OUT_OF_DISTRIBUTION

        # Tier 1: TRIVIAL - high mean_logit indicates all models likely correct
        if self.mean_logit > thresholds.trivial_mean_logit:
            return InferenceTier.TRIVIAL

        # Tier 3: UNCERTAIN - need ensemble for quality
        # Low mean_logit = hard query, low margin = model can't distinguish
        if (
            self.mean_logit < thresholds.uncertain_mean_logit
            or self.logit_margin < thresholds.uncertain_logit_margin
        ):
            return InferenceTier.UNCERTAIN

        # Tier 2: DISCRIMINATIVE - GNN routing adds value
        return InferenceTier.DISCRIMINATIVE


@dataclass
class TieredRoutingResult:
    """Result of tiered routing decision."""

    tier: InferenceTier
    signals: DifficultySignals
    selected_model: str
    selected_config: TierConfig
    embedding_latency_ms: float = 0.0
    classification_latency_ms: float = 0.0
    rationale: str = ""

    @property
    def total_routing_latency_ms(self) -> float:
        return self.embedding_latency_ms + self.classification_latency_ms


class TieredRouter:
    """Intelligent tiered routing based on GNN difficulty prediction.

    This router extends EB1Router with automatic tier classification,
    selecting the optimal inference strategy based on query difficulty.
    """

    def __init__(
        self,
        gnn_model: torch.nn.Module,
        llm_embeddings: torch.Tensor,
        llm_names: list[str],
        llm_metadata: dict[str, dict],
        llm_model_ids: dict[str, str] | None = None,
        thresholds: TieredThresholds | None = None,
        tier_configs: dict[InferenceTier, TierConfig] | None = None,
        device: str = "cpu",
        embed_dim: int = 512,
        graph_cache: dict | None = None,
        inference_weights: tuple[float, float, float] = (0.7, 0.2, 0.1),
    ):
        """Initialize tiered router.

        Args:
            gnn_model: Trained GNN model for routing predictions.
            llm_embeddings: Pre-computed LLM description embeddings.
            llm_names: List of LLM names in order.
            llm_metadata: Metadata (cost, latency) for each LLM.
            llm_model_ids: Mapping from display names to Ember model IDs.
            thresholds: Tier classification thresholds.
            tier_configs: Configuration for each tier.
            device: Device for inference ("cpu" or "cuda").
            embed_dim: Embedding dimension for queries.
            graph_cache: Training graph tensors for full GNN inference.
                Required for correct routing quality. Contains query/task
                embeddings, edge features, and train/valid masks from the
                training dataset.
            inference_weights: Tuple of (perf, cost, latency) preference
                weights for edge feature augmentation.
        """
        if graph_cache is None:
            raise ValueError(
                "graph_cache is required for TieredRouter. The full GNN "
                "forward pass requires the training graph for message "
                "passing context. Load it from training_graph_cache.pkl."
            )

        self.gnn_model = gnn_model
        self.llm_embeddings = llm_embeddings.to(device)
        self.llm_names = llm_names
        self.llm_metadata = llm_metadata
        # Mapping from display names to Ember model IDs for API calls
        self.llm_model_ids = llm_model_ids or {
            name: name.lower().replace(" ", "-") for name in llm_names
        }
        self.thresholds = thresholds or TieredThresholds()
        self.tier_configs = tier_configs or DEFAULT_TIER_CONFIGS.copy()
        self.device = device
        self.embed_dim = embed_dim
        self.num_llms = len(llm_names)

        # Store training graph tensors for full GNN forward pass
        self._graph_query_embs = (
            torch.from_numpy(graph_cache["query_embeddings"])
            .float()
            .to(device)
        )
        self._graph_task_embs = (
            torch.from_numpy(graph_cache["task_embeddings"])
            .float()
            .to(device)
        )
        self._graph_combined_edge = (
            torch.from_numpy(graph_cache["combined_edge"])
            .float()
            .to(device)
        )
        self._graph_source_nodes = (
            torch.from_numpy(graph_cache["source_nodes"])
            .long()
            .to(device)
        )
        self._graph_dest_nodes = (
            torch.from_numpy(graph_cache["dest_nodes"])
            .long()
            .to(device)
        )
        train_mask = (
            torch.from_numpy(graph_cache["train_mask"])
            .bool()
            .to(device)
        )
        valid_mask = (
            torch.from_numpy(graph_cache["valid_mask"])
            .bool()
            .to(device)
        )
        self._graph_edge_can_see = train_mask | valid_mask
        self._graph_num_queries = graph_cache["num_queries"]
        self._inference_weights = inference_weights

        # Initialize embedding client
        self._embedding_client = None

        # Set model to eval mode
        self.gnn_model.eval()

        num_visible = int(self._graph_edge_can_see.sum())
        logger.info(
            "TieredRouter initialized with full graph context: "
            "%d queries, %d visible edges, %d LLMs",
            self._graph_num_queries,
            num_visible,
            self.num_llms,
        )

    def _get_embedding_client(self):
        """Lazy-load embedding client."""
        if self._embedding_client is None:
            from embergraphrouter.inference.eb1_router import EmbeddingClient

            self._embedding_client = EmbeddingClient(
                model="text-embedding-3-large",
                dimensions=self.embed_dim,
            )
        return self._embedding_client

    @classmethod
    def from_registry(
        cls,
        registry_name: str,
        registry_root: str | Path | None = None,
        thresholds: TieredThresholds | None = None,
        tier_configs: dict[InferenceTier, TierConfig] | None = None,
        device: str = "cpu",
    ) -> "TieredRouter":
        """Load tiered router from registry.

        Automatically loads calibrated thresholds from the registry's
        thresholds.yaml file when available. Each GNN checkpoint produces a
        specific logit distribution, so thresholds must be calibrated per model.

        Args:
            registry_name: Name of registry entry.
            registry_root: Root directory for registry.
            thresholds: Custom tier thresholds. If None, loads from registry's
                thresholds.yaml (recommended). Pass explicit thresholds only for
                experimentation.
            tier_configs: Custom tier configurations.
            device: Device for inference.

        Returns:
            Initialized TieredRouter with calibrated thresholds.
        """
        import json
        import pickle

        import yaml

        if registry_root is None:
            registry_root = Path(__file__).parent.parent / "registry"
        registry_root = Path(registry_root)
        registry_path = registry_root / registry_name

        if not registry_path.exists():
            raise FileNotFoundError(f"Registry not found: {registry_path}")

        # Load calibrated thresholds from registry if not explicitly provided
        if thresholds is None:
            thresholds = TieredThresholds.from_registry(registry_path, fallback_to_defaults=True)

        # Load config
        config_path = registry_path / "config.yaml"
        with open(config_path) as f:
            config_dict = yaml.safe_load(f)

        embed_dim = config_dict.get("embed_dim", 512)
        hidden_dim = config_dict.get("hidden_dim", 384)
        num_gat_layers = config_dict.get("num_gat_layers", 3)
        num_gat_heads = config_dict.get("num_gat_heads", 4)
        dropout = config_dict.get("dropout", 0.4)
        edge_feature_dim = config_dict.get("edge_feature_dim", 6)

        # Load LLM embeddings
        emb_path = registry_path / "llm_embeddings.pkl"
        with open(emb_path, "rb") as f:
            llm_embeddings_np = pickle.load(f)
        llm_embeddings = torch.from_numpy(llm_embeddings_np).float()

        # Load LLM descriptions
        desc_path = registry_path / "llm_descriptions.json"
        with open(desc_path) as f:
            llm_descriptions = json.load(f)

        # NOTE: llm_names will be derived from graph_cache["llm_order"] below
        # for guaranteed ordering consistency with training. The
        # llm_descriptions dict is only used for metadata/model_ids.
        llm_metadata = {
            name: {
                "cost": desc.get("cost", 0.0),
                "latency": desc.get("latency", 0.0),
            }
            for name, desc in llm_descriptions.items()
        }
        # Build display_name -> model_id mapping for Ember API calls
        llm_model_ids = {
            name: desc.get("model", name.lower().replace(" ", "-"))
            for name, desc in llm_descriptions.items()
        }

        # Initialize model
        try:
            from embergraphrouter.model.graph_nn import GATEncoderDecoderNet
        except ModuleNotFoundError:
            from model.graph_nn import GATEncoderDecoderNet

        # Load weights first to detect architecture from state_dict
        model_path = registry_path / "model.pth"
        state_dict = torch.load(model_path, map_location=device, weights_only=True)
        use_score_head = "score_head.weight" in state_dict

        model = GATEncoderDecoderNet(
            query_feature_dim=embed_dim,
            llm_feature_dim=embed_dim,
            hidden_features=hidden_dim,
            in_edges=edge_feature_dim,
            num_gnn_layers=num_gat_layers,
            num_attention_heads=num_gat_heads,
            dropout=dropout,
            hidden_multiplier=2,
            use_contrastive=True,
            use_score_head=use_score_head,
            use_difficulty_head=True,
            use_uncertainty_head=False,
        )

        # Load weights (strict=False to ignore deprecated keys)
        model.load_state_dict(state_dict, strict=False)
        model.to(device)

        # Load training graph cache (REQUIRED for full GNN inference)
        cache_path = registry_path / "training_graph_cache.pkl"
        if not cache_path.exists():
            raise FileNotFoundError(
                f"Training graph cache not found: {cache_path}. "
                "The full GNN forward pass requires the training graph "
                "for message passing context. Run "
                "package_inference_cache.py to create it."
            )
        with open(cache_path, "rb") as f:
            graph_cache = pickle.load(f)

        # CRITICAL: Use llm_order from training graph as authoritative source.
        # This ensures GNN output indices map to correct model names.
        # The graph_cache stores the exact order used during training.
        if "llm_order" not in graph_cache:
            raise ValueError(
                f"Training graph cache missing 'llm_order' key: {cache_path}. "
                "This is required to ensure GNN indices map to correct models."
            )
        llm_names = graph_cache["llm_order"]

        # Validate consistency with llm_descriptions.json
        json_names = set(llm_descriptions.keys())
        cache_names = set(llm_names)
        if json_names != cache_names:
            missing_in_json = cache_names - json_names
            missing_in_cache = json_names - cache_names
            raise ValueError(
                f"LLM name mismatch between training cache and descriptions. "
                f"In cache but not JSON: {missing_in_json}. "
                f"In JSON but not cache: {missing_in_cache}."
            )

        # Load inference weights from config
        iw = config_dict.get("inference_weights", {})
        inference_weights = (
            iw.get("perf", 0.7),
            iw.get("cost", 0.2),
            iw.get("latency", 0.1),
        )

        return cls(
            gnn_model=model,
            llm_embeddings=llm_embeddings,
            llm_names=llm_names,
            llm_metadata=llm_metadata,
            llm_model_ids=llm_model_ids,
            thresholds=thresholds,
            tier_configs=tier_configs,
            device=device,
            embed_dim=embed_dim,
            graph_cache=graph_cache,
            inference_weights=inference_weights,
        )

    def extract_contrastive_similarity(
        self, query_embedding: torch.Tensor
    ) -> float | None:
        """Extract max contrastive similarity between query and LLMs.

        This provides a principled OOD detection signal based on the contrastive
        training objective. The contrastive encoder is trained via InfoNCE to pull
        in-distribution queries toward their best-performing LLM. Therefore:
        - High max similarity → Query aligns with at least one LLM (in-distribution)
        - Low max similarity → Query doesn't align with any LLM (out-of-distribution)

        Args:
            query_embedding: Query embedding tensor (1, embed_dim).

        Returns:
            Max cosine similarity to any LLM in contrastive space, or None if
            the model doesn't have a contrastive encoder.
        """
        # Check if model has contrastive encoder
        if not hasattr(self.gnn_model, "contrastive_encoder"):
            return None
        if self.gnn_model.contrastive_encoder is None:
            return None

        with torch.no_grad():
            # Get aligned representations from FeatureAlign
            # This is the same projection used during GNN forward pass
            task_emb = query_embedding  # Task = query for single-query inference
            x_initial = self.gnn_model.model_align(
                task_emb, query_embedding, self.llm_embeddings
            )

            # First row is query, rest are LLMs
            query_repr = x_initial[:1]
            llm_repr = x_initial[1:]

            # Project through contrastive encoder's projections
            contrastive_encoder = self.gnn_model.contrastive_encoder
            q_proj = F.normalize(contrastive_encoder.query_proj(query_repr), dim=-1)
            l_proj = F.normalize(contrastive_encoder.llm_proj(llm_repr), dim=-1)

            # Compute similarities (cosine similarity due to normalization)
            similarities = torch.matmul(q_proj, l_proj.T).squeeze(0)
            max_similarity = float(similarities.max().item())

        return max_similarity

    def extract_signals(self, embedding: np.ndarray | torch.Tensor) -> DifficultySignals:
        """Extract difficulty signals from query embedding.

        Args:
            embedding: Query embedding vector.

        Returns:
            DifficultySignals with all classification inputs.
        """
        if isinstance(embedding, np.ndarray):
            embedding = torch.from_numpy(embedding).float()
        embedding = embedding.to(self.device)

        if embedding.dim() == 1:
            embedding = embedding.unsqueeze(0)

        with torch.no_grad():
            logits = self._forward(embedding)

        logits_np = logits.squeeze(0).cpu().numpy()
        probs = F.softmax(logits, dim=-1).squeeze(0).cpu().numpy()

        # Compute signals
        mean_logit = float(np.mean(logits_np))
        logit_std = float(np.std(logits_np))

        # Top-2 for margin computation
        sorted_indices = np.argsort(logits_np)[::-1]
        top1_idx = int(sorted_indices[0])
        top2_idx = int(sorted_indices[1])
        logit_margin = float(logits_np[top1_idx] - logits_np[top2_idx])

        max_prob = float(probs[top1_idx])

        # Entropy (for reference, less useful due to softmax compression)
        entropy = -float(np.sum(probs * np.log(probs + 1e-8)))
        max_entropy = np.log(self.num_llms)
        normalized_entropy = entropy / max_entropy

        # Extract contrastive similarity for principled OOD detection
        max_contrastive_similarity = self.extract_contrastive_similarity(embedding)

        return DifficultySignals(
            mean_logit=mean_logit,
            max_prob=max_prob,
            logit_margin=logit_margin,
            logit_std=logit_std,
            entropy=normalized_entropy,
            top1_model=self.llm_names[top1_idx],
            top1_idx=top1_idx,
            top2_model=self.llm_names[top2_idx],
            top2_idx=top2_idx,
            all_probs={name: float(probs[i]) for i, name in enumerate(self.llm_names)},
            all_logits={name: float(logits_np[i]) for i, name in enumerate(self.llm_names)},
            max_contrastive_similarity=max_contrastive_similarity,
        )

    def _forward(self, query_embedding: torch.Tensor) -> torch.Tensor:
        """Run full GNN forward pass with training graph context.

        Appends the new query to the training graph and runs the complete
        GATEncoderDecoderNet forward pass. The GAT layers use train+valid
        edges for message passing, enriching LLM representations with
        learned quality signals from thousands of training queries.

        Args:
            query_embedding: Query embedding tensor of shape (1, embed_dim).

        Returns:
            Logits tensor of shape (1, num_llms).
        """
        new_query_idx = self._graph_num_queries
        num_queries_extended = new_query_idx + 1

        # Extend query/task embeddings with new query
        extended_query = torch.cat(
            [self._graph_query_embs, query_embedding], dim=0
        )
        extended_task = torch.cat(
            [self._graph_task_embs, query_embedding], dim=0
        )

        # Build new edges: new_query -> all LLMs
        new_src = torch.full(
            (self.num_llms,), new_query_idx,
            dtype=torch.long, device=self.device,
        )
        new_dst = torch.arange(
            self.num_llms, dtype=torch.long, device=self.device
        )

        # Extended edge index with dest offset for node indexing.
        # Nodes: [query_0..query_N, new_query, llm_0..llm_K]
        # LLM nodes start at index num_queries_extended.
        ext_src = torch.cat([self._graph_source_nodes, new_src])
        ext_dst = torch.cat([
            self._graph_dest_nodes + num_queries_extended,
            new_dst + num_queries_extended,
        ])
        edge_index = torch.stack([ext_src, ext_dst])

        # Edge features: [metrics, preferences] per graph_inference convention
        new_edge_metrics = torch.zeros(
            self.num_llms, 3, device=self.device
        )
        ext_combined = torch.cat(
            [self._graph_combined_edge, new_edge_metrics], dim=0
        )
        wp, wc, wl = self._inference_weights
        pref = torch.tensor(
            [[wp, wc, wl]], device=self.device
        ).expand(len(ext_combined), -1)
        edge_weight = torch.cat([ext_combined, pref], dim=1)

        # Masks: see train+valid edges, predict only new query's edges
        num_new = self.num_llms
        edge_can_see = torch.cat([
            self._graph_edge_can_see,
            torch.zeros(num_new, dtype=torch.bool, device=self.device),
        ])
        edge_mask = torch.cat([
            torch.zeros(
                len(ext_src) - num_new,
                dtype=torch.bool, device=self.device,
            ),
            torch.ones(num_new, dtype=torch.bool, device=self.device),
        ])

        logits = self.gnn_model(
            task_features=extended_task,
            query_features=extended_query,
            llm_features=self.llm_embeddings,
            edge_index=edge_index,
            edge_mask=edge_mask,
            edge_can_see=edge_can_see,
            edge_weight=edge_weight,
        )

        if logits.dim() == 1:
            logits = logits.unsqueeze(0)
        return logits

    def classify(self, query: str) -> TieredRoutingResult:
        """Classify query into tier and return routing decision.

        Args:
            query: Query text to classify.

        Returns:
            TieredRoutingResult with tier, signals, and selected model.
        """
        # Get embedding
        client = self._get_embedding_client()
        embed_start = time.perf_counter()
        embedding, _ = client.embed(query)
        embed_latency = (time.perf_counter() - embed_start) * 1000

        # Extract signals and classify
        classify_start = time.perf_counter()
        signals = self.extract_signals(embedding)
        tier = signals.classify_tier(self.thresholds)
        classify_latency = (time.perf_counter() - classify_start) * 1000

        config = self.tier_configs[tier]

        # Determine selected model
        if config.use_ensemble:
            selected_model = f"eb1:{config.ensemble_plan}"
        elif tier == InferenceTier.DISCRIMINATIVE:
            # Use GNN's top pick for discriminative tier
            selected_model = signals.top1_model
        elif tier == InferenceTier.OUT_OF_DISTRIBUTION:
            # OOD: sample ensemble vs top-tier based on ood_ensemble_probability
            # This enables exploration while maintaining safety
            if (
                self.thresholds.ood_ensemble_probability > 0
                and random.random() < self.thresholds.ood_ensemble_probability
            ):
                # Use ensemble for exploration
                uncertain_config = self.tier_configs.get(InferenceTier.UNCERTAIN)
                if uncertain_config and uncertain_config.use_ensemble:
                    selected_model = f"eb1:{uncertain_config.ensemble_plan}"
                else:
                    selected_model = config.model
            else:
                # Default: use top-tier model for safety
                selected_model = config.model
        else:
            selected_model = config.model

        # Build rationale
        rationale = self._build_rationale(tier, signals)

        return TieredRoutingResult(
            tier=tier,
            signals=signals,
            selected_model=selected_model,
            selected_config=config,
            embedding_latency_ms=embed_latency,
            classification_latency_ms=classify_latency,
            rationale=rationale,
        )

    def _build_rationale(self, tier: InferenceTier, signals: DifficultySignals) -> str:
        """Build human-readable rationale for tier selection."""
        if tier == InferenceTier.OUT_OF_DISTRIBUTION:
            # Determine which OOD detection method was used
            if (
                signals.max_contrastive_similarity is not None
                and self.thresholds.ood_contrastive_threshold > 0
                and signals.max_contrastive_similarity < self.thresholds.ood_contrastive_threshold
            ):
                return (
                    f"OUT_OF_DISTRIBUTION: contrastive_similarity={signals.max_contrastive_similarity:.3f} < "
                    f"{self.thresholds.ood_contrastive_threshold:.3f}. "
                    f"Query has low alignment to all LLMs in contrastive space; "
                    f"routing to top-tier model."
                )
            else:
                return (
                    f"OUT_OF_DISTRIBUTION: confidence={signals.max_prob:.3f} < "
                    f"{self.thresholds.ood_max_confidence}, "
                    f"logit_std={signals.logit_std:.3f} < "
                    f"{self.thresholds.ood_max_logit_std}. "
                    f"Query outside training distribution; routing to top-tier model."
                )
        elif tier == InferenceTier.TRIVIAL:
            return (
                f"TRIVIAL: mean_logit={signals.mean_logit:.3f} > "
                f"{self.thresholds.trivial_mean_logit}. "
                f"All models likely correct; routing to cheapest."
            )
        elif tier == InferenceTier.UNCERTAIN:
            reasons = []
            if signals.mean_logit < self.thresholds.uncertain_mean_logit:
                reasons.append(
                    f"mean_logit={signals.mean_logit:.3f} < {self.thresholds.uncertain_mean_logit}"
                )
            if signals.logit_margin < self.thresholds.uncertain_logit_margin:
                reasons.append(
                    f"margin={signals.logit_margin:.3f} < {self.thresholds.uncertain_logit_margin}"
                )
            return f"UNCERTAIN: {'; '.join(reasons)}. Using ensemble for quality."
        else:
            return (
                f"DISCRIMINATIVE: GNN predicts {signals.top1_model} "
                f"(prob={signals.max_prob:.3f}, margin={signals.logit_margin:.3f}). "
                f"Routing adds value."
            )

    def route(self, query: str) -> TieredRoutingResult:
        """Route query to appropriate tier and model.

        This is the main entry point for tiered routing. It:
        1. Embeds the query
        2. Extracts difficulty signals from GNN
        3. Classifies into tier
        4. Returns routing decision with rationale

        Args:
            query: Query text to route.

        Returns:
            TieredRoutingResult with full routing decision.
        """
        return self.classify(query)

    def route_and_execute(
        self,
        query: str,
        *,
        system: str | None = None,
        shared_options: Mapping[str, object] | None = None,
    ) -> "ChatResponse":
        """Route query and execute inference.

        This method performs end-to-end tiered inference:
        1. Classify query into tier
        2. Route to appropriate model/ensemble
        3. Execute inference
        4. Return response with routing metadata

        Args:
            query: Query text.
            system: Optional system message.
            shared_options: Options passed to model invocation.

        Returns:
            ChatResponse with routing metadata attached.
        """
        routing_result = self.classify(query)

        if routing_result.selected_config.use_ensemble:
            response = self._execute_ensemble(
                query,
                routing_result,
                system=system,
                shared_options=shared_options,
            )
        else:
            response = self._execute_single_model(
                query,
                routing_result,
                system=system,
                shared_options=shared_options,
            )

        # Attach routing metadata
        self._attach_routing_metadata(response, routing_result)
        return response

    def _execute_single_model(
        self,
        query: str,
        routing_result: TieredRoutingResult,
        *,
        system: str | None = None,
        shared_options: Mapping[str, object] | None = None,
    ) -> "ChatResponse":
        """Execute inference with a single model."""
        from ember.api.models import models

        options = dict(shared_options or {})
        options.update(routing_result.selected_config.model_params)

        if system:
            options["system"] = system

        # Resolve display name to Ember model ID
        display_name = routing_result.selected_model
        model_id = self.llm_model_ids.get(display_name, display_name)

        return models.response(model_id, query, **options)

    def _execute_ensemble(
        self,
        query: str,
        routing_result: TieredRoutingResult,
        *,
        system: str | None = None,
        shared_options: Mapping[str, object] | None = None,
    ) -> "ChatResponse":
        """Execute inference with EB1 ensemble."""
        from ember.takumi.eb1 import (
            DEFAULT_PLAN,
            DEFAULT_PRO_HIGH_PLAN,
            DEFAULT_PRO_PLAN,
            run_plan,
        )

        plan_name = routing_result.selected_config.ensemble_plan
        if plan_name == "DEFAULT_PRO_HIGH_PLAN":
            plan = DEFAULT_PRO_HIGH_PLAN
        elif plan_name == "DEFAULT_PRO_PLAN":
            plan = DEFAULT_PRO_PLAN
        else:
            plan = DEFAULT_PLAN

        return run_plan(
            plan,
            query,
            system=system,
            shared_options=shared_options,
        )

    def _attach_routing_metadata(
        self,
        response: object,
        routing_result: TieredRoutingResult,
    ) -> None:
        """Attach tiered routing metadata to response.

        Works with both Response wrapper and ChatResponse dataclass.
        """
        # Get the underlying ChatResponse if wrapped
        if hasattr(response, "chat_response"):
            # Response wrapper from ember.api.models
            chat_response = response.chat_response
        elif hasattr(response, "raw_output"):
            # Direct ChatResponse
            chat_response = response
        else:
            # Can't attach metadata to unknown type
            return

        raw_output = chat_response.raw_output
        if not isinstance(raw_output, dict):
            raw_output = {"original": raw_output}

        signals = routing_result.signals
        signals_dict = {
            "mean_logit": signals.mean_logit,
            "max_prob": signals.max_prob,
            "logit_margin": signals.logit_margin,
            "logit_std": signals.logit_std,
            "entropy": signals.entropy,
            "top1_model": signals.top1_model,
            "top2_model": signals.top2_model,
        }
        # Include contrastive similarity if available
        if signals.max_contrastive_similarity is not None:
            signals_dict["max_contrastive_similarity"] = signals.max_contrastive_similarity

        thresholds_dict = {
            "trivial_mean_logit": self.thresholds.trivial_mean_logit,
            "uncertain_mean_logit": self.thresholds.uncertain_mean_logit,
            "uncertain_logit_margin": self.thresholds.uncertain_logit_margin,
        }
        # Include contrastive threshold if enabled
        if self.thresholds.ood_contrastive_threshold > 0:
            thresholds_dict["ood_contrastive_threshold"] = (
                self.thresholds.ood_contrastive_threshold
            )

        raw_output["tiered_routing"] = {
            "tier": routing_result.tier.name,
            "selected_model": routing_result.selected_model,
            "rationale": routing_result.rationale,
            "signals": signals_dict,
            "thresholds": thresholds_dict,
            "latency_ms": {
                "embedding": routing_result.embedding_latency_ms,
                "classification": routing_result.classification_latency_ms,
                "total_routing": routing_result.total_routing_latency_ms,
            },
        }

        # Update the underlying dataclass directly
        chat_response.raw_output = raw_output


def calibrate_thresholds_from_data(
    router: TieredRouter,
    queries: Sequence[str],
    ground_truth_effects: np.ndarray,
    target_trivial_pct: float = 0.30,
    target_uncertain_pct: float = 0.15,
) -> TieredThresholds:
    """Calibrate thresholds from labeled data.

    This function finds optimal thresholds that achieve target tier distributions
    while maximizing routing quality. It's useful for domain-specific tuning.

    Args:
        router: Initialized TieredRouter.
        queries: List of calibration queries.
        ground_truth_effects: Matrix of shape (n_queries, n_llms) with true effects.
        target_trivial_pct: Target percentage of queries in TRIVIAL tier.
        target_uncertain_pct: Target percentage of queries in UNCERTAIN tier.

    Returns:
        Calibrated TieredThresholds.
    """
    # Extract signals for all queries
    signals_list = []
    for query in queries:
        client = router._get_embedding_client()
        embedding, _ = client.embed(query)
        signals = router.extract_signals(embedding)
        signals_list.append(signals)

    # Compute ground truth difficulty metrics
    n_correct = (ground_truth_effects >= 0.5).sum(axis=1)
    pct_correct = n_correct / ground_truth_effects.shape[1]
    perf_range = ground_truth_effects.max(axis=1) - ground_truth_effects.min(axis=1)

    # Extract signal arrays
    mean_logits = np.array([s.mean_logit for s in signals_list])
    logit_margins = np.array([s.logit_margin for s in signals_list])

    # Find trivial threshold: queries where all models correct
    trivial_mask = pct_correct >= 0.9
    if trivial_mask.sum() > 0:
        trivial_mean_logit = float(
            np.percentile(mean_logits[trivial_mask], 100 - target_trivial_pct * 100)
        )
    else:
        trivial_mean_logit = 0.5

    # Find uncertain threshold: queries with high perf_range (routing matters)
    uncertain_mask = (pct_correct <= 0.3) | (perf_range >= 0.5)
    if uncertain_mask.sum() > 0:
        uncertain_mean_logit = float(
            np.percentile(mean_logits[uncertain_mask], target_uncertain_pct * 100)
        )
        uncertain_logit_margin = float(
            np.percentile(logit_margins[uncertain_mask], target_uncertain_pct * 100)
        )
    else:
        uncertain_mean_logit = -0.5
        uncertain_logit_margin = 0.1

    return TieredThresholds(
        trivial_mean_logit=trivial_mean_logit,
        uncertain_mean_logit=uncertain_mean_logit,
        uncertain_logit_margin=uncertain_logit_margin,
    )


__all__ = [
    "TieredRouter",
    "TieredThresholds",
    "TierConfig",
    "InferenceTier",
    "DifficultySignals",
    "TieredRoutingResult",
    "calibrate_thresholds_from_data",
    "DEFAULT_TIER_CONFIGS",
]
