"""Multi-Head MLP Router baseline.

This architecture uses a shared query encoder with separate prediction heads
for each model in the routing catalog.

Architecture:
    Query -> [Shared Encoder] -> query_repr
                                    |
                    +---------------+---------------+
                    |               |               |
                 [Head_0]        [Head_1]   ...  [Head_N]
                    |               |               |
                 score_0         score_1   ...   score_N

Advantages:
    - Simple and interpretable
    - Fast inference (single forward pass)
    - Each head can specialize to model-specific patterns
    - Easy to add/remove models (add/remove heads)

Disadvantages:
    - Parameters scale O(N) with number of models
    - No parameter sharing across model heads
    - No explicit model feature usage (heads are index-based)

When to use:
    - Small model catalog (< 20 models)
    - Stable model catalog (infrequent additions)
    - Need for interpretability (which head fires for which query type)
    - Latency-sensitive applications
"""

from __future__ import annotations

from typing import Optional

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

from .base import BaseRouter, ResidualBlock, RouterConfig, SEBlock


class MultiHeadMLPRouter(BaseRouter):
    """Multi-headed MLP router with shared query encoder.

    Each model gets its own prediction head, allowing specialization while
    sharing the query encoding computation.
    """

    def __init__(self, config: RouterConfig) -> None:
        """Initialize Multi-Head MLP router.

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

        # Shared query encoder
        self.query_encoder = self._build_encoder(
            input_dim=config.query_dim,
            hidden_dim=config.hidden_dim,
            num_layers=config.num_layers,
            dropout=config.dropout,
            use_layer_norm=config.use_layer_norm,
            use_residual=config.use_residual,
        )

        # Optional SE block for query feature importance
        self.query_se = None
        if config.use_se_block:
            self.query_se = SEBlock(config.hidden_dim, reduction=4)

        # Per-model prediction heads
        # Each head: hidden_dim -> 1 (scalar score)
        self.heads = nn.ModuleList([
            self._build_head(config.hidden_dim, config.dropout)
            for _ in range(config.num_models)
        ])

        # Learnable temperature for calibration
        self.log_temperature = nn.Parameter(
            torch.log(torch.tensor(config.temperature))
        )

        # Model embeddings are not used in this architecture,
        # but we store them for interface compatibility
        self._cached_model_repr: Optional[torch.Tensor] = None

    def _build_encoder(
        self,
        input_dim: int,
        hidden_dim: int,
        num_layers: int,
        dropout: float,
        use_layer_norm: bool,
        use_residual: bool,
    ) -> nn.Module:
        """Build the shared query encoder.

        Args:
            input_dim: Input embedding dimension.
            hidden_dim: Hidden layer dimension.
            num_layers: Number of layers.
            dropout: Dropout probability.
            use_layer_norm: Whether to use layer normalization.
            use_residual: Whether to use residual connections.

        Returns:
            Encoder module.
        """
        layers = []

        # Input projection
        layers.append(nn.Linear(input_dim, hidden_dim))
        if use_layer_norm:
            layers.append(nn.LayerNorm(hidden_dim))
        layers.append(nn.GELU())
        layers.append(nn.Dropout(dropout))

        # Residual blocks or simple MLPs
        for _ in range(num_layers - 1):
            if use_residual:
                layers.append(
                    ResidualBlock(
                        dim=hidden_dim,
                        hidden_mult=2.0,
                        dropout=dropout,
                        use_layer_norm=use_layer_norm,
                    )
                )
            else:
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                if use_layer_norm:
                    layers.append(nn.LayerNorm(hidden_dim))
                layers.append(nn.GELU())
                layers.append(nn.Dropout(dropout))

        return nn.Sequential(*layers)

    def _build_head(self, hidden_dim: int, dropout: float) -> nn.Module:
        """Build a single prediction head.

        Args:
            hidden_dim: Input dimension from encoder.
            dropout: Dropout probability.

        Returns:
            Head module outputting scalar score.
        """
        return nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )

    def encode_query(self, query_emb: torch.Tensor) -> torch.Tensor:
        """Encode query embeddings through shared encoder.

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

        Returns:
            Encoded query representation, shape (batch, hidden_dim).
        """
        query_repr = self.query_encoder(query_emb)

        if self.query_se is not None:
            query_repr = self.query_se(query_repr)

        return query_repr

    def encode_models(self, model_embs: torch.Tensor) -> torch.Tensor:
        """Encode model embeddings (pass-through for this architecture).

        In Multi-Head MLP, model identities are implicit in the head indices,
        not in learned representations. This method exists for interface
        compatibility.

        Args:
            model_embs: Model embeddings, shape (num_models, model_dim).
                       Not actually used for routing, but cached for analysis.

        Returns:
            Pass-through of input (for interface compatibility).
        """
        self._cached_model_repr = model_embs
        return model_embs

    def compute_scores(
        self, query_repr: torch.Tensor, model_repr: torch.Tensor
    ) -> torch.Tensor:
        """Compute routing scores via per-model heads.

        Args:
            query_repr: Encoded queries, shape (batch, hidden_dim).
            model_repr: Model embeddings (unused, for interface compatibility).

        Returns:
            Routing scores, shape (batch, num_models).
        """
        # Apply each head to get per-model scores
        scores_list = [head(query_repr) for head in self.heads]

        # Stack: list of (batch, 1) -> (batch, num_models)
        scores = torch.cat(scores_list, dim=-1)

        # Apply learned temperature
        temperature = torch.exp(self.log_temperature)
        scores = scores / temperature

        return scores

    def get_head_weights(self) -> list[torch.Tensor]:
        """Get the weight matrices from each head for interpretability.

        Returns:
            List of weight tensors from the first layer of each head.
        """
        return [head[0].weight.detach() for head in self.heads]

    def get_head_activations(
        self, query_emb: torch.Tensor
    ) -> dict[int, torch.Tensor]:
        """Get intermediate activations from each head for analysis.

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

        Returns:
            Dictionary mapping head index to intermediate activations.
        """
        query_repr = self.encode_query(query_emb)
        activations = {}
        for i, head in enumerate(self.heads):
            # Get activation after first layer
            intermediate = head[0](query_repr)
            intermediate = head[1](intermediate)  # GELU
            activations[i] = intermediate.detach()
        return activations


class MultiHeadMLPWithModelFeatures(BaseRouter):
    """Multi-Head MLP variant that incorporates model features.

    This variant uses model embeddings to condition the prediction heads,
    combining the simplicity of multi-head with some model awareness.

    Architecture:
        Query -> [Shared Encoder] -> query_repr (384-dim)
        Model -> [Model Encoder]  -> model_repr (384-dim)

        For each model i:
            score_i = Head_i(query_repr) + bias_from(model_repr_i)

    This allows the model features to influence routing while keeping
    separate heads for per-model specialization.
    """

    def __init__(self, config: RouterConfig) -> None:
        """Initialize Multi-Head MLP with model features.

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

        # Shared query encoder
        self.query_encoder = nn.Sequential(
            nn.Linear(config.query_dim, config.hidden_dim),
            nn.LayerNorm(config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            ResidualBlock(config.hidden_dim, dropout=config.dropout),
        )

        # Model encoder (projects model features to hidden space)
        self.model_encoder = nn.Sequential(
            nn.Linear(config.model_dim, config.hidden_dim),
            nn.LayerNorm(config.hidden_dim),
            nn.GELU(),
        )

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

        # Bias projection from model features
        self.model_bias = nn.Linear(config.hidden_dim, 1)

        # Learnable combination weight
        self.alpha = nn.Parameter(torch.tensor(0.5))

    def encode_query(self, query_emb: torch.Tensor) -> torch.Tensor:
        """Encode query embeddings."""
        return self.query_encoder(query_emb)

    def encode_models(self, model_embs: torch.Tensor) -> torch.Tensor:
        """Encode model embeddings."""
        return self.model_encoder(model_embs)

    def compute_scores(
        self, query_repr: torch.Tensor, model_repr: torch.Tensor
    ) -> torch.Tensor:
        """Compute scores combining head outputs and model biases.

        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).
        """
        batch_size = query_repr.shape[0]
        num_models = model_repr.shape[0]

        # Head-based scores
        head_scores = torch.cat(
            [head(query_repr) for head in self.heads], dim=-1
        )  # (batch, num_models)

        # Model feature-based bias
        model_biases = self.model_bias(model_repr).squeeze(-1)  # (num_models,)
        model_biases = model_biases.unsqueeze(0).expand(batch_size, -1)

        # Combine with learned weight
        alpha = torch.sigmoid(self.alpha)
        scores = alpha * head_scores + (1 - alpha) * model_biases

        return scores
