"""Base classes and interfaces for router baselines.

This module defines the common interface that all router architectures implement,
enabling fair comparison and interchangeable usage in training/evaluation pipelines.

Design Principles:
    - Explicit configuration via RouterConfig dataclass
    - Consistent forward signature: (query_emb, model_embs) -> scores
    - Unified difficulty signals for downstream tiered routing
    - No magic: all behavior is explicit and testable
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


class RouterArchitecture(Enum):
    """Enumeration of available router architectures."""

    GNN = "gnn"
    MULTI_HEAD_MLP = "multi_head_mlp"
    TWO_TOWER = "two_tower"


@dataclass
class RouterConfig:
    """Configuration for router architectures.

    This config is shared across all baseline implementations. Each architecture
    uses the subset of parameters relevant to its design.

    Attributes:
        architecture: Which router architecture to use.
        query_dim: Dimension of input query embeddings.
        model_dim: Dimension of input model/LLM embeddings.
        hidden_dim: Hidden layer dimension for encoders.
        num_models: Number of LLMs in the routing catalog.
        num_layers: Number of encoder layers (MLP depth or GNN layers).
        num_heads: Number of attention heads (for GNN/Two-Tower attention).
        dropout: Dropout probability for regularization.
        temperature: Temperature for score calibration.
        use_layer_norm: Whether to use layer normalization.
        use_residual: Whether to use residual connections.
        use_se_block: Whether to use Squeeze-and-Excitation for feature importance.
    """

    architecture: RouterArchitecture = RouterArchitecture.TWO_TOWER
    query_dim: int = 512
    model_dim: int = 512
    hidden_dim: int = 384
    num_models: int = 7
    num_layers: int = 2
    num_heads: int = 4
    dropout: float = 0.1
    temperature: float = 1.0
    use_layer_norm: bool = True
    use_residual: bool = True
    use_se_block: bool = True

    def __post_init__(self) -> None:
        """Validate configuration values."""
        if self.query_dim <= 0:
            raise ValueError(f"query_dim must be positive, got {self.query_dim}")
        if self.model_dim <= 0:
            raise ValueError(f"model_dim must be positive, got {self.model_dim}")
        if self.hidden_dim <= 0:
            raise ValueError(f"hidden_dim must be positive, got {self.hidden_dim}")
        if self.num_models <= 0:
            raise ValueError(f"num_models must be positive, got {self.num_models}")
        if not 0 <= self.dropout < 1:
            raise ValueError(f"dropout must be in [0, 1), got {self.dropout}")
        if self.temperature <= 0:
            raise ValueError(f"temperature must be positive, got {self.temperature}")


@dataclass
class DifficultySignals:
    """Signals extracted from router predictions for tiered routing.

    These signals enable downstream decision-making about query difficulty:
    - Trivial queries: All models likely correct, use cheapest
    - Discriminative queries: Clear winner, use GNN top pick
    - Uncertain queries: High disagreement, use ensemble

    Attributes:
        mean_logit: Mean of raw logits (higher = easier query).
        max_prob: Maximum probability after softmax (confidence in top choice).
        logit_margin: Gap between top-2 logits (decisiveness).
        entropy: Entropy of probability distribution (uncertainty).
        top_k_agreement: Fraction of top-k models that are "good" choices.
    """

    mean_logit: float
    max_prob: float
    logit_margin: float
    entropy: float
    top_k_agreement: float = 0.0

    @classmethod
    def from_scores(
        cls, scores: torch.Tensor, k: int = 3
    ) -> "DifficultySignals":
        """Compute difficulty signals from raw routing scores.

        Args:
            scores: Raw routing scores, shape (num_models,) or (batch, num_models).
            k: Number of top models for agreement calculation.

        Returns:
            DifficultySignals instance with computed values.
        """
        if scores.dim() == 1:
            scores = scores.unsqueeze(0)

        # Compute probabilities
        probs = F.softmax(scores, dim=-1)

        # Mean logit (averaged over batch if batched)
        mean_logit = scores.mean(dim=-1).mean().item()

        # Max probability (confidence)
        max_prob = probs.max(dim=-1).values.mean().item()

        # Logit margin (top-1 minus top-2)
        sorted_scores, _ = scores.sort(dim=-1, descending=True)
        logit_margin = (sorted_scores[:, 0] - sorted_scores[:, 1]).mean().item()

        # Entropy (normalized by log(num_models))
        num_models = scores.shape[-1]
        entropy_raw = -(probs * (probs + 1e-10).log()).sum(dim=-1)
        max_entropy = torch.log(torch.tensor(num_models, dtype=scores.dtype))
        entropy = (entropy_raw / max_entropy).mean().item()

        # Top-k agreement: fraction of probability mass in top-k
        top_k_probs, _ = probs.topk(min(k, num_models), dim=-1)
        top_k_agreement = top_k_probs.sum(dim=-1).mean().item()

        return cls(
            mean_logit=mean_logit,
            max_prob=max_prob,
            logit_margin=logit_margin,
            entropy=entropy,
            top_k_agreement=top_k_agreement,
        )


@dataclass
class RouterOutput:
    """Output from router forward pass.

    Attributes:
        scores: Raw routing scores, shape (batch, num_models).
        query_repr: Encoded query representation (for analysis/debugging).
        model_repr: Encoded model representations (if computed).
        signals: Difficulty signals for tiered routing.
    """

    scores: torch.Tensor
    query_repr: Optional[torch.Tensor] = None
    model_repr: Optional[torch.Tensor] = None
    signals: Optional[DifficultySignals] = None


class BaseRouter(nn.Module, ABC):
    """Abstract base class for router architectures.

    All router implementations must inherit from this class and implement
    the abstract methods. This ensures consistent interface across architectures.

    The forward pass signature is:
        forward(query_emb, model_embs, return_signals=False) -> RouterOutput

    For inference, use:
        route(query_emb, model_embs) -> (best_model_idx, confidence, signals)
    """

    def __init__(self, config: RouterConfig) -> None:
        """Initialize the router with configuration.

        Args:
            config: RouterConfig instance with architecture parameters.
        """
        super().__init__()
        self.config = config

    @abstractmethod
    def encode_query(self, query_emb: torch.Tensor) -> torch.Tensor:
        """Encode query embeddings into routing representation.

        Args:
            query_emb: Query embeddings, shape (batch, query_dim).

        Returns:
            Encoded query representation, shape (batch, hidden_dim).
        """
        pass

    @abstractmethod
    def encode_models(self, model_embs: torch.Tensor) -> torch.Tensor:
        """Encode model/LLM embeddings into routing representation.

        Args:
            model_embs: Model embeddings, shape (num_models, model_dim).

        Returns:
            Encoded model representations, shape (num_models, hidden_dim).
        """
        pass

    @abstractmethod
    def compute_scores(
        self, query_repr: torch.Tensor, model_repr: torch.Tensor
    ) -> torch.Tensor:
        """Compute routing scores from encoded representations.

        Args:
            query_repr: Encoded queries, shape (batch, hidden_dim).
            model_repr: Encoded models, shape (num_models, hidden_dim).

        Returns:
            Routing scores, shape (batch, num_models).
        """
        pass

    def forward(
        self,
        query_emb: torch.Tensor,
        model_embs: torch.Tensor,
        return_signals: bool = False,
    ) -> RouterOutput:
        """Forward pass computing routing scores.

        Args:
            query_emb: Query embeddings, shape (batch, query_dim).
            model_embs: Model embeddings, shape (num_models, model_dim).
            return_signals: Whether to compute difficulty signals.

        Returns:
            RouterOutput with scores and optional representations/signals.
        """
        query_repr = self.encode_query(query_emb)
        model_repr = self.encode_models(model_embs)
        scores = self.compute_scores(query_repr, model_repr)

        signals = None
        if return_signals:
            signals = DifficultySignals.from_scores(scores)

        return RouterOutput(
            scores=scores,
            query_repr=query_repr,
            model_repr=model_repr,
            signals=signals,
        )

    def route(
        self, query_emb: torch.Tensor, model_embs: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, DifficultySignals]:
        """Route queries to best models with confidence and signals.

        Convenience method for inference that returns actionable routing decisions.

        Args:
            query_emb: Query embeddings, shape (batch, query_dim).
            model_embs: Model embeddings, shape (num_models, model_dim).

        Returns:
            Tuple of:
                - best_model_idx: Index of best model per query, shape (batch,).
                - confidence: Probability of best model, shape (batch,).
                - signals: DifficultySignals for tiered routing.
        """
        output = self.forward(query_emb, model_embs, return_signals=True)
        probs = F.softmax(output.scores / self.config.temperature, dim=-1)
        confidence, best_model_idx = probs.max(dim=-1)
        return best_model_idx, confidence, output.signals

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

    def get_architecture_name(self) -> str:
        """Return human-readable architecture name."""
        return self.config.architecture.value


class ResidualBlock(nn.Module):
    """Pre-norm residual MLP block.

    Architecture: LayerNorm -> Linear -> Activation -> Dropout -> Linear -> Add

    This follows the modern transformer pattern of pre-normalization for
    more stable training dynamics.
    """

    def __init__(
        self,
        dim: int,
        hidden_mult: float = 4.0,
        dropout: float = 0.1,
        use_layer_norm: bool = True,
    ) -> None:
        """Initialize residual block.

        Args:
            dim: Input and output dimension.
            hidden_mult: Multiplier for hidden dimension. Default 4.0.
            dropout: Dropout probability. Default 0.1.
            use_layer_norm: Whether to use layer normalization. Default True.
        """
        super().__init__()
        hidden_dim = int(dim * hidden_mult)

        self.norm = nn.LayerNorm(dim) if use_layer_norm else nn.Identity()
        self.fc1 = nn.Linear(dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, dim)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with residual connection.

        Args:
            x: Input tensor, shape (..., dim).

        Returns:
            Output tensor with same shape as input.
        """
        residual = x
        x = self.norm(x)
        x = self.fc1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x + residual


class SEBlock(nn.Module):
    """Squeeze-and-Excitation block for feature importance learning.

    Learns channel-wise (feature-wise) attention weights to emphasize
    important features and suppress less relevant ones.

    Reference: Hu et al., "Squeeze-and-Excitation Networks" (CVPR 2018)
    """

    def __init__(self, dim: int, reduction: int = 4) -> None:
        """Initialize SE block.

        Args:
            dim: Input dimension.
            reduction: Reduction factor for bottleneck. Default 4.
        """
        super().__init__()
        bottleneck = max(dim // reduction, 8)
        self.fc1 = nn.Linear(dim, bottleneck)
        self.fc2 = nn.Linear(bottleneck, dim)
        self.activation = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass computing and applying attention weights.

        Args:
            x: Input tensor, shape (..., dim).

        Returns:
            Reweighted tensor with same shape as input.
        """
        # Squeeze: global average (already per-sample, just compute weights)
        weights = self.fc1(x)
        weights = self.activation(weights)
        weights = self.fc2(weights)
        weights = torch.sigmoid(weights)

        # Excitation: element-wise reweighting
        return x * weights
