"""Cross-Attention architecture for LLM routing.

This module implements a simpler, more natural architecture for the query→LLM
scoring problem. Instead of graph message passing, queries directly attend
to LLM embeddings to compute routing scores.

Key insight: LLM routing is fundamentally a scoring/matching problem, not
a graph reasoning problem. Cross-attention matches this structure better.
"""

from __future__ import annotations

import math

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


class CrossAttentionRouter(nn.Module):
    """Cross-attention based router for query→LLM scoring.

    Architecture:
        1. Project query and LLM embeddings to hidden dimension
        2. Multi-head cross-attention: query attends to all LLMs
        3. Combine with edge features (preferences)
        4. MLP decoder to produce per-LLM scores

    This is simpler than GAT because:
        - No message passing between queries (they're independent)
        - No iterative graph convolutions
        - Direct query→LLM attention matches the problem structure

    Args:
        query_feature_dim: Query embedding dimension (e.g., 4096 for Qwen3)
        llm_feature_dim: LLM embedding dimension
        hidden_dim: Hidden dimension for projections
        num_llms: Number of LLMs to route to
        edge_dim: Edge feature dimension (preferences)
        num_heads: Number of attention heads
        num_layers: Number of cross-attention layers
        dropout: Dropout rate
        use_difficulty_head: Whether to predict query difficulty
    """

    def __init__(
        self,
        query_feature_dim: int,
        llm_feature_dim: int,
        hidden_dim: int = 512,
        num_llms: int = 7,
        edge_dim: int = 6,
        num_heads: int = 8,
        num_layers: int = 2,
        dropout: float = 0.1,
        use_difficulty_head: bool = False,
        use_uncertainty_head: bool = False,
    ) -> None:
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_llms = num_llms
        self.num_heads = num_heads
        self.use_difficulty_head = use_difficulty_head
        self.use_uncertainty_head = use_uncertainty_head

        # Input projections
        self.query_proj = nn.Sequential(
            nn.Linear(query_feature_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )

        self.llm_proj = nn.Sequential(
            nn.Linear(llm_feature_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )

        # Edge feature projection
        self.edge_proj = nn.Sequential(
            nn.Linear(edge_dim, hidden_dim // 4),
            nn.GELU(),
            nn.Linear(hidden_dim // 4, hidden_dim // 4),
        )

        # Cross-attention layers
        self.cross_attn_layers = nn.ModuleList()
        self.ffn_layers = nn.ModuleList()
        self.layer_norms1 = nn.ModuleList()
        self.layer_norms2 = nn.ModuleList()

        for _ in range(num_layers):
            self.cross_attn_layers.append(
                nn.MultiheadAttention(
                    embed_dim=hidden_dim,
                    num_heads=num_heads,
                    dropout=dropout,
                    batch_first=True,
                )
            )
            self.ffn_layers.append(
                nn.Sequential(
                    nn.Linear(hidden_dim, hidden_dim * 4),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim * 4, hidden_dim),
                    nn.Dropout(dropout),
                )
            )
            self.layer_norms1.append(nn.LayerNorm(hidden_dim))
            self.layer_norms2.append(nn.LayerNorm(hidden_dim))

        # Score decoder: combines query repr, llm repr, edge features
        decoder_input_dim = hidden_dim * 2 + hidden_dim // 4
        self.score_decoder = nn.Sequential(
            nn.Linear(decoder_input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 1),
        )

        # Optional difficulty head
        if use_difficulty_head:
            self.difficulty_head = nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim // 2),
                nn.GELU(),
                nn.Linear(hidden_dim // 2, 1),
                nn.Sigmoid(),
            )

        # Optional uncertainty head
        if use_uncertainty_head:
            self.uncertainty_head = nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim // 2),
                nn.GELU(),
                nn.Linear(hidden_dim // 2, 1),
                nn.Sigmoid(),
            )

        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights with Xavier/Glorot."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(
        self,
        query_emb: torch.Tensor,
        llm_emb: torch.Tensor,
        edge_features: torch.Tensor,
        edge_index: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Forward pass.

        Args:
            query_emb: Query embeddings (batch_size, query_dim)
            llm_emb: LLM embeddings (num_llms, llm_dim)
            edge_features: Edge features (batch_size * num_llms, edge_dim)
            edge_index: Not used, kept for compatibility with GNN interface

        Returns:
            Dictionary with 'scores' and optionally 'difficulty', 'uncertainty'
        """
        batch_size = query_emb.shape[0]
        num_llms = llm_emb.shape[0]

        # Project inputs
        query_h = self.query_proj(query_emb)  # (batch, hidden)
        llm_h = self.llm_proj(llm_emb)  # (num_llms, hidden)

        # Expand LLM embeddings for batch
        llm_h = llm_h.unsqueeze(0).expand(batch_size, -1, -1)  # (batch, num_llms, hidden)

        # Cross-attention: query attends to LLMs
        query_h = query_h.unsqueeze(1)  # (batch, 1, hidden)

        for cross_attn, ffn, ln1, ln2 in zip(
            self.cross_attn_layers, self.ffn_layers,
            self.layer_norms1, self.layer_norms2
        ):
            # Cross-attention with residual
            attn_out, _ = cross_attn(query_h, llm_h, llm_h)
            query_h = ln1(query_h + attn_out)

            # FFN with residual
            query_h = ln2(query_h + ffn(query_h))

        query_h = query_h.squeeze(1)  # (batch, hidden)

        # Compute scores for each query-LLM pair
        # Reshape edge features: (batch * num_llms, edge_dim) -> (batch, num_llms, edge_dim)
        edge_h = self.edge_proj(edge_features)  # (batch * num_llms, hidden//4)
        edge_h = edge_h.view(batch_size, num_llms, -1)  # (batch, num_llms, hidden//4)

        # Get per-LLM representations
        llm_h_squeezed = self.llm_proj(llm_emb)  # (num_llms, hidden)
        llm_h_expanded = llm_h_squeezed.unsqueeze(0).expand(batch_size, -1, -1)  # (batch, num_llms, hidden)

        # Expand query representation for each LLM
        query_h_expanded = query_h.unsqueeze(1).expand(-1, num_llms, -1)  # (batch, num_llms, hidden)

        # Concatenate: [query_repr, llm_repr, edge_features]
        combined = torch.cat([query_h_expanded, llm_h_expanded, edge_h], dim=-1)

        # Decode scores
        scores = self.score_decoder(combined).squeeze(-1)  # (batch, num_llms)

        result = {"scores": scores}

        # Optional heads
        if self.use_difficulty_head:
            result["difficulty"] = self.difficulty_head(query_h).squeeze(-1)

        if self.use_uncertainty_head:
            result["uncertainty"] = self.uncertainty_head(query_h).squeeze(-1)

        return result


class CrossAttentionEncoderDecoderNet(nn.Module):
    """Wrapper to match the GATEncoderDecoderNet interface.

    This adapter makes CrossAttentionRouter compatible with the existing
    training loop by matching the forward signature and output format.
    """

    def __init__(
        self,
        query_feature_dim: int,
        llm_feature_dim: int,
        hidden_features: int,
        in_edges: int,
        dropout: float = 0.1,
        num_gnn_layers: int = 2,  # Maps to num_layers
        use_difficulty_head: bool = False,
        difficulty_head_mode: str = "regression",
        num_attention_heads: int = 8,
        attention_dropout: float = 0.1,
        hidden_multiplier: int = 2,
        query_mlp_layers: int = 1,
        llm_mlp_layers: int = 1,
        input_dropout: float = 0.0,
        use_residual: bool = False,
        use_swiglu: bool = False,
        use_variational: bool = False,
        vib_beta: float = 0.001,
        use_gated_skip: bool = False,
        use_query_moe: bool = False,
        moe_num_experts: int = 4,
        moe_top_k: int = 2,
        use_contrastive: bool = False,
        contrastive_weight: float = 0.1,
        use_uniformity_loss: bool = False,
        uniformity_weight: float = 0.05,
        use_uncertainty_head: bool = False,
        **kwargs,  # Absorb any extra args for compatibility
    ) -> None:
        super().__init__()

        self.use_difficulty_head = use_difficulty_head
        self.use_contrastive = use_contrastive
        self.contrastive_weight = contrastive_weight
        self.use_uniformity_loss = use_uniformity_loss
        self.uniformity_weight = uniformity_weight

        self.router = CrossAttentionRouter(
            query_feature_dim=query_feature_dim,
            llm_feature_dim=llm_feature_dim,
            hidden_dim=hidden_features,
            num_llms=7,  # Will be set dynamically
            edge_dim=in_edges,
            num_heads=num_attention_heads,
            num_layers=num_gnn_layers,
            dropout=dropout,
            use_difficulty_head=use_difficulty_head,
            use_uncertainty_head=use_uncertainty_head,
        )

        self.input_dropout = nn.Dropout(input_dropout) if input_dropout > 0 else nn.Identity()

    def forward(
        self,
        query_x: torch.Tensor,
        llm_x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
        """Forward pass matching GATEncoderDecoderNet interface.

        Args:
            query_x: Query embeddings (num_queries, query_dim)
            llm_x: LLM embeddings (num_llms, llm_dim)
            edge_index: Edge indices (2, num_edges) - used to infer batch structure
            edge_attr: Edge features (num_edges, edge_dim)

        Returns:
            Tuple of (scores, difficulty, uncertainty)
        """
        # Apply input dropout
        query_x = self.input_dropout(query_x)
        llm_x = self.input_dropout(llm_x)

        # Forward through router
        result = self.router(query_x, llm_x, edge_attr, edge_index)

        scores = result["scores"]
        difficulty = result.get("difficulty")
        uncertainty = result.get("uncertainty")

        return scores, difficulty, uncertainty

    def get_query_embeddings(self, query_x: torch.Tensor) -> torch.Tensor:
        """Get query embeddings for contrastive loss."""
        return self.router.query_proj(query_x)
