"""GNN architectures for LLM routing prediction.

Contains the encoder-decoder network architectures (GCN and GAT variants) that
combine feature alignment, message passing, and edge score prediction.
"""

from __future__ import annotations

import logging

import torch
from torch import nn
from torch.nn import functional as F
from torch_geometric.nn import GATv2Conv, GeneralConv

from embergraphrouter.model.layers import (
    ContrastiveQueryEncoder,
    FeatureAlign,
    GatedIdentitySkip,
    PreferenceHypernetwork,
    QueryMoE,
)

logger = logging.getLogger(__name__)


class EncoderDecoderNet(nn.Module):
    """GNN encoder-decoder for edge score prediction.

    Architecture:
        1. FeatureAlign: Projects query/task/LLM features to common space
        2. EdgeMLP: Transforms edge features (preferences + metrics)
        3. GeneralConv x2-3: Message passing with residual connections
        4. Edge prediction: Dot product between query and LLM representations

    Note: difficulty_head was removed. Query difficulty is now computed from
    mean logit, which provides 6x better correlation (0.31 vs 0.05).
    See scripts/eval_difficulty_aware.py for the mean logit analysis.
    """

    def __init__(
        self,
        query_feature_dim: int,
        llm_feature_dim: int,
        hidden_features: int,
        in_edges: int,
        dropout: float = 0.0,
        num_gnn_layers: int = 2,
        hidden_multiplier: int = 2,
        query_mlp_layers: int = 1,
        input_dropout: float = 0.0,
        # Legacy parameters (ignored, kept for checkpoint compatibility)
        use_difficulty_head: bool = False,
        difficulty_head_mode: str = "regression",
    ) -> None:
        # Warn if deprecated parameters are used
        if use_difficulty_head:
            logger.warning(
                "use_difficulty_head is deprecated. Query difficulty is now derived "
                "from mean logit, which provides 6x better correlation (0.31 vs 0.05)."
            )
        super().__init__()
        self.in_edges = in_edges
        self.num_gnn_layers = num_gnn_layers
        self.hidden_multiplier = hidden_multiplier
        self.model_align = FeatureAlign(
            query_feature_dim,
            llm_feature_dim,
            hidden_features,
            query_mlp_layers=query_mlp_layers,
            input_dropout=input_dropout,
        )

        channels = hidden_features * hidden_multiplier
        align_output_dim = hidden_features * 2  # FeatureAlign outputs 2*hidden_features

        # Projection layer to bridge FeatureAlign output to GNN input
        # FeatureAlign outputs 2*hidden_features, GNN expects hidden_features*hidden_multiplier
        if align_output_dim != channels:
            self.input_projection = nn.Linear(align_output_dim, channels)
        else:
            self.input_projection = None

        # Build GNN layers dynamically based on num_gnn_layers
        self.conv_layers = nn.ModuleList()
        self.bn_layers = nn.ModuleList()
        for _ in range(num_gnn_layers):
            self.conv_layers.append(
                GeneralConv(
                    in_channels=channels,
                    out_channels=channels,
                    in_edge_channels=in_edges,
                )
            )
            self.bn_layers.append(nn.BatchNorm1d(channels))

        # 2-layer edge MLP for richer edge feature transformation
        self.edge_mlp = nn.Sequential(
            nn.Linear(in_edges, in_edges * 2),
            nn.ReLU(),
            nn.Linear(in_edges * 2, in_edges),
        )
        self.dropout = nn.Dropout(dropout)

        # Note: difficulty_head has been removed.
        # Query difficulty is now computed from mean logit (0.31 correlation vs 0.05).
        # See scripts/eval_difficulty_aware.py for the mean logit analysis.

    def forward(
        self,
        task_features: torch.Tensor,
        query_features: torch.Tensor,
        llm_features: torch.Tensor,
        edge_index: torch.Tensor,
        edge_mask: torch.Tensor,
        edge_can_see: torch.Tensor,
        edge_weight: torch.Tensor,
        best_llm_idx: torch.Tensor | None = None,
        preferences: torch.Tensor | None = None,
        # Legacy parameters (ignored, kept for API compatibility)
        return_difficulty: bool = False,
    ) -> torch.Tensor:
        """Forward pass predicting scores for masked edges.

        Args:
            task_features: Task description embeddings (num_queries, dim).
            query_features: Query embeddings (num_queries, dim).
            llm_features: LLM description embeddings (num_llms, dim).
            edge_index: Edge connectivity (2, num_edges).
            edge_mask: Boolean mask for edges to predict.
            edge_can_see: Boolean mask for edges visible during message passing.
            edge_weight: Edge features including preferences (num_edges, in_edges).
            best_llm_idx: Unused in GCN, kept for API compatibility with GAT.
            preferences: Unused in GCN, kept for API compatibility with GAT.
            return_difficulty: DEPRECATED - ignored, difficulty from mean logit.

        Returns:
            edge_logits: Predicted scores for masked edges (num_masked_edges,).
            Note: For difficulty detection, compute mean(logits) per query.
            See DifficultyDetection in scripts/eval_difficulty_aware.py.
        """
        # Note: best_llm_idx and preferences are unused in GCN (no V6+ modules)
        del best_llm_idx, preferences, return_difficulty

        edge_index_visible = edge_index[:, edge_can_see]
        edge_index_predict = edge_index[:, edge_mask]
        edge_weight_visible = edge_weight[edge_can_see]
        edge_weight_visible = self.edge_mlp(
            edge_weight_visible.reshape(-1, self.in_edges)
        ).relu_()

        x_initial = self.model_align(task_features, query_features, llm_features)

        # Project to GNN input dimension if needed
        if self.input_projection is not None:
            x_initial = self.input_projection(x_initial)

        # Run through all GNN layers with residual connections
        x = x_initial
        for i, (conv, bn) in enumerate(zip(self.conv_layers, self.bn_layers, strict=False)):
            x_new = conv(x, edge_index_visible, edge_attr=edge_weight_visible)
            x_new = bn(x_new)
            x_new = F.relu(x_new)
            x_new = self.dropout(x_new)
            # Residual connection after first layer
            if i > 0:
                x_new = x_new + x
            x = x_new

        x_out = x

        # Edge prediction via dot product
        query_repr = x_initial[edge_index_predict[0]]
        llm_repr = x_out[edge_index_predict[1]]
        edge_logits = (query_repr * llm_repr).mean(dim=-1)

        return edge_logits


class GATEncoderDecoderNet(nn.Module):
    """GAT-based encoder-decoder for edge score prediction.

    Architecture identical to EncoderDecoderNet but uses Graph Attention Networks
    (GATv2Conv) instead of GeneralConv for message passing.

    Key differences from GCN:
    - Multi-head attention mechanism learns importance of neighbors
    - Each head computes attention coefficients independently
    - More expressive than fixed message passing in GeneralConv
    - Slightly higher computational cost (4 heads x hidden_dim)

    Optional architectural improvements (V6+):
    - Gated Identity Skip: Learnable alpha * x_initial + (1-alpha) * x_out blending
    - Query MoE: Route queries through specialized expert sub-networks
    - Preference Hypernetwork: Generate routing params from preferences
    - Contrastive Pretraining: Pull queries toward best LLM, push from worst

    Expected improvements:
    - Better handling of heterogeneous query-LLM relationships
    - More interpretable (attention weights show routing decisions)
    - +2-4% frac_of_oracle in preliminary analysis
    """

    def __init__(
        self,
        query_feature_dim: int,
        llm_feature_dim: int,
        hidden_features: int,
        in_edges: int,
        dropout: float = 0.0,
        num_gnn_layers: int = 2,
        num_attention_heads: int = 4,
        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,
        # V6+ architectural improvements
        use_gated_skip: bool = False,
        use_query_moe: bool = False,
        moe_num_experts: int = 4,
        moe_top_k: int = 2,
        use_hypernetwork: bool = False,
        use_contrastive: bool = False,
        contrastive_weight: float = 0.1,
        use_uniformity_loss: bool = False,
        uniformity_weight: float = 0.05,
        use_score_head: bool = False,
        # Legacy parameters (ignored, kept for checkpoint compatibility)
        use_difficulty_head: bool = False,
        difficulty_head_mode: str = "regression",
        use_uncertainty_head: bool = False,
    ) -> None:
        """Initialize GAT encoder-decoder.

        Args:
            query_feature_dim: Query embedding dimension.
            llm_feature_dim: LLM embedding dimension.
            hidden_features: Hidden dimension for GNN layers.
            in_edges: Edge feature dimension (preferences + metrics).
            dropout: Dropout rate for node features.
            num_gnn_layers: Number of GAT layers (2-3 recommended).
            num_attention_heads: Number of attention heads per GAT layer.
            attention_dropout: Dropout rate for attention coefficients.
            hidden_multiplier: Multiplier for hidden channels (default 2).
            query_mlp_layers: Number of MLP layers for query preprocessing.
            llm_mlp_layers: Number of MLP layers for LLM preprocessing.
            input_dropout: Dropout on input embeddings (data augmentation).
            use_residual: Use residual MLP blocks instead of plain MLP.
            use_swiglu: Use SwiGLU activation (requires use_residual=True).
            use_variational: Apply variational information bottleneck.
            vib_beta: Weight for KL divergence loss when use_variational=True.
            use_gated_skip: Enable learnable x_initial/x_out interpolation.
            use_query_moe: Enable Mixture of Experts for query processing.
            moe_num_experts: Number of MoE experts (default 4).
            moe_top_k: Number of experts activated per query (default 2).
            use_hypernetwork: Generate routing params from preferences.
            use_contrastive: Enable contrastive learning for embeddings.
            contrastive_weight: Weight for contrastive loss (default 0.1).
            use_uniformity_loss: Enable uniformity loss in contrastive learning.
            uniformity_weight: Weight for uniformity loss (default 0.05).
            use_score_head: Replace fixed mean pooling with a learned Linear(channels, 1)
                scoring function. Discovers which embedding dimensions carry routing
                signal instead of averaging all equally (which collapses logits by CLT).
            use_difficulty_head: DEPRECATED - difficulty derived from mean logit.
            difficulty_head_mode: DEPRECATED - no longer used.
            use_uncertainty_head: DEPRECATED - no longer used.
        """
        # Warn if deprecated parameters are used
        if use_difficulty_head:
            logger.warning(
                "use_difficulty_head is deprecated. Query difficulty is now derived "
                "from mean logit, which provides 6x better correlation (0.31 vs 0.05)."
            )
        if use_uncertainty_head:
            logger.warning(
                "use_uncertainty_head is deprecated and will be ignored."
            )
        super().__init__()
        self.in_edges = in_edges
        self.num_gnn_layers = num_gnn_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_multiplier = hidden_multiplier
        self.use_variational = use_variational
        self.vib_beta = vib_beta

        # V6+ feature flags
        self.use_gated_skip = use_gated_skip
        self.use_query_moe = use_query_moe
        self.use_hypernetwork = use_hypernetwork
        self.use_contrastive = use_contrastive
        self.contrastive_weight = contrastive_weight

        self.model_align = FeatureAlign(
            query_feature_dim,
            llm_feature_dim,
            hidden_features,
            query_mlp_layers=query_mlp_layers,
            llm_mlp_layers=llm_mlp_layers,
            input_dropout=input_dropout,
            use_residual=use_residual,
            use_swiglu=use_swiglu,
            use_variational=use_variational,
        )
        channels = hidden_features * hidden_multiplier
        align_output_dim = hidden_features * 2  # FeatureAlign outputs 2*hidden_features

        # Projection layer to bridge FeatureAlign output to GNN input
        # FeatureAlign outputs 2*hidden_features, GNN expects hidden_features*hidden_multiplier
        if align_output_dim != channels:
            self.input_projection = nn.Linear(align_output_dim, channels)
        else:
            self.input_projection = None

        # Build GAT layers with multi-head attention
        self.conv_layers = nn.ModuleList()
        self.bn_layers = nn.ModuleList()

        for _ in range(num_gnn_layers):
            self.conv_layers.append(
                GATv2Conv(
                    in_channels=channels,
                    out_channels=channels,
                    edge_dim=in_edges,
                    heads=num_attention_heads,
                    concat=False,  # Average heads instead of concatenating
                    dropout=attention_dropout,
                    add_self_loops=False,  # Bipartite graph, no self-loops
                )
            )
            self.bn_layers.append(nn.BatchNorm1d(channels))

        # Edge MLP (same as GCN)
        self.edge_mlp = nn.Sequential(
            nn.Linear(in_edges, in_edges * 2),
            nn.ReLU(),
            nn.Linear(in_edges * 2, in_edges),
        )
        self.dropout = nn.Dropout(dropout)

        # V6+ Gated Identity Skip for query representation
        if use_gated_skip:
            self.query_gated_skip = GatedIdentitySkip(channels, per_feature=True)
            self.llm_gated_skip = GatedIdentitySkip(channels, per_feature=True)
        else:
            self.query_gated_skip = None
            self.llm_gated_skip = None

        # V6+ Mixture of Experts for query processing
        if use_query_moe:
            self.query_moe = QueryMoE(
                dim=channels,
                num_experts=moe_num_experts,
                top_k=moe_top_k,
                dropout=dropout,
            )
        else:
            self.query_moe = None

        # V6+ Preference Hypernetwork
        if use_hypernetwork:
            # Generates weights for final routing projection
            # Input: 3 preference values (perf, cost, latency)
            # Output: projection from channels -> 1 (score)
            self.preference_hypernet = PreferenceHypernetwork(
                pref_dim=3,
                hidden_dim=64,
                output_weight_shape=(channels, 1),
            )
        else:
            self.preference_hypernet = None

        # V6+ Contrastive query encoder with optional uniformity loss
        if use_contrastive:
            self.contrastive_encoder = ContrastiveQueryEncoder(
                query_dim=channels,
                llm_dim=channels,
                projection_dim=128,
                temperature=0.07,
                use_uniformity=use_uniformity_loss,
                uniformity_weight=uniformity_weight,
            )
        else:
            self.contrastive_encoder = None

        # Learned score head: replaces fixed mean pooling with a linear projection
        # that discovers which embedding dimensions carry routing signal.
        # Mean pooling averages all channels equally, which by CLT washes out
        # discriminative signal and produces near-identical logits for all LLMs.
        #
        # Initialize as mean pooling (weight = 1/channels, bias = 0) so the
        # model starts with v14-equivalent behavior. This avoids gradient
        # starvation: random small weights (~0.002) choke the gradient flowing
        # back to the GNN, preventing it from learning discriminative
        # representations. Starting from mean pooling, the head learns to
        # deviate toward informative dimensions over training.
        if use_score_head:
            self.score_head = nn.Linear(channels, 1, bias=True)
            nn.init.constant_(self.score_head.weight, 1.0 / channels)
            nn.init.zeros_(self.score_head.bias)
        else:
            self.score_head = None

        # Note: difficulty_head and uncertainty_head have been removed.
        # Query difficulty is now computed from mean logit (0.31 correlation vs 0.05).
        # See scripts/eval_difficulty_aware.py for the mean logit analysis.

    def forward(
        self,
        task_features: torch.Tensor,
        query_features: torch.Tensor,
        llm_features: torch.Tensor,
        edge_index: torch.Tensor,
        edge_mask: torch.Tensor,
        edge_can_see: torch.Tensor,
        edge_weight: torch.Tensor,
        best_llm_idx: torch.Tensor | None = None,
        preferences: torch.Tensor | None = None,
        # Legacy parameters (ignored, kept for API compatibility)
        return_difficulty: bool = False,
        return_uncertainty: bool = False,
    ) -> torch.Tensor:
        """Forward pass with GAT message passing and optional V6+ features.

        Args:
            task_features: Task description embeddings (num_queries, dim).
            query_features: Query embeddings (num_queries, dim).
            llm_features: LLM description embeddings (num_llms, dim).
            edge_index: Edge connectivity (2, num_edges).
            edge_mask: Boolean mask for edges to predict.
            edge_can_see: Boolean mask for edges visible during message passing.
            edge_weight: Edge features including preferences (num_edges, in_edges).
            best_llm_idx: Best LLM index per query for contrastive learning (optional).
            preferences: Preference vector [perf, cost, latency] for hypernetwork (optional).
            return_difficulty: DEPRECATED - ignored, difficulty from mean logit.
            return_uncertainty: DEPRECATED - ignored.

        Returns:
            edge_logits: Predicted scores for masked edges (num_masked_edges,).
            Note: For difficulty detection, compute mean(logits) per query.
            See DifficultyDetection in scripts/eval_difficulty_aware.py.
        """
        edge_index_visible = edge_index[:, edge_can_see]
        edge_index_predict = edge_index[:, edge_mask]
        edge_weight_visible = edge_weight[edge_can_see]
        edge_weight_visible = self.edge_mlp(
            edge_weight_visible.reshape(-1, self.in_edges)
        ).relu_()

        x_initial = self.model_align(task_features, query_features, llm_features)
        num_queries = task_features.size(0)
        num_llms = llm_features.size(0)

        # Project to GNN input dimension if needed
        if self.input_projection is not None:
            x_initial = self.input_projection(x_initial)

        # V6+ Apply MoE to query features before GNN
        if self.query_moe is not None:
            query_initial = x_initial[:num_queries]
            query_initial = self.query_moe(query_initial)
            x_initial = torch.cat([query_initial, x_initial[num_queries:]], dim=0)

        # V6+ Contrastive + uniformity on trainable aligned embeddings.
        # This shapes the router's internal representation (FeatureAlign / MoE / projection),
        # not the frozen input embeddings from the dataset.
        if self.contrastive_encoder is not None:
            self.contrastive_encoder.clear_losses()
            if self.training:
                if best_llm_idx is None:
                    raise ValueError(
                        "best_llm_idx must be provided during training when use_contrastive is enabled"
                    )
                if best_llm_idx.dim() != 1 or best_llm_idx.numel() != num_queries:
                    raise ValueError(
                        "best_llm_idx must have shape (num_queries,) for contrastive learning: "
                        f"got shape={tuple(best_llm_idx.shape)} num_queries={num_queries}"
                    )
                if best_llm_idx.dtype != torch.long:
                    raise ValueError(
                        "best_llm_idx must be torch.long for contrastive learning: "
                        f"got dtype={best_llm_idx.dtype}"
                    )

                predict_query_indices = edge_index_predict[0].unique()
                if predict_query_indices.numel() > 0:
                    query_repr = x_initial.index_select(0, predict_query_indices)
                    llm_repr = x_initial.narrow(0, num_queries, num_llms)
                    best_llm_idx_batch = best_llm_idx.index_select(0, predict_query_indices)

                    _ = self.contrastive_encoder(query_repr, llm_repr, best_llm_idx_batch)

        # Run through all GAT layers with residual connections
        x = x_initial
        for i, (conv, bn) in enumerate(zip(self.conv_layers, self.bn_layers, strict=False)):
            # GAT forward returns (node_features, attention_weights)
            # We only use node_features; attention weights could be logged for interpretability
            x_new = conv(x, edge_index_visible, edge_attr=edge_weight_visible)
            x_new = bn(x_new)
            x_new = F.relu(x_new)
            x_new = self.dropout(x_new)
            # Residual connection after first layer
            if i > 0:
                x_new = x_new + x
            x = x_new

        x_out = x

        # V6+ Apply gated skip connections for edge prediction
        if self.query_gated_skip is not None:
            query_repr_initial = x_initial[edge_index_predict[0]]
            query_repr_out = x_out[edge_index_predict[0]]
            query_repr = self.query_gated_skip(query_repr_initial, query_repr_out)
        else:
            # Default: use x_initial for queries (preserves query-intrinsic info)
            query_repr = x_initial[edge_index_predict[0]]

        if self.llm_gated_skip is not None:
            llm_repr_initial = x_initial[edge_index_predict[1]]
            llm_repr_out = x_out[edge_index_predict[1]]
            llm_repr = self.llm_gated_skip(llm_repr_initial, llm_repr_out)
        else:
            # Default: use x_out for LLMs (benefits from GNN context)
            llm_repr = x_out[edge_index_predict[1]]

        # Combined edge representation - used for both quality and uncertainty
        # This ensures uncertainty reflects confidence in the SAME features used for prediction
        edge_repr = query_repr * llm_repr  # (num_edges, channels)

        # V6+ Hypernetwork: generate preference-conditioned routing weights
        if self.preference_hypernet is not None and preferences is not None:
            weight, bias = self.preference_hypernet(preferences)
            # Apply generated weights: (num_edges, channels) @ (channels, 1) + bias
            edge_logits = (edge_repr @ weight).squeeze(-1) + bias.squeeze(-1)
        elif self.score_head is not None:
            # v15: Learned scoring -- discovers which dimensions carry routing signal
            edge_logits = self.score_head(edge_repr).squeeze(-1)
        else:
            # Legacy: mean pooling (backward compatible with v14 weights)
            edge_logits = edge_repr.mean(dim=-1)

        return edge_logits

    def get_auxiliary_losses(self) -> dict[str, torch.Tensor]:
        """Collect auxiliary losses from V6+ modules for training.

        Returns:
            Dictionary of auxiliary losses to add to main loss:
            - 'moe_balance': Load balancing loss for MoE (if enabled)
            - 'contrastive': Contrastive learning loss (if enabled)
            - 'uniformity': Uniformity loss for embedding spread (if enabled)
            - 'vib_kl': KL divergence from variational bottleneck (if enabled)
        """
        losses = {}

        # MoE load balancing loss
        if self.query_moe is not None and self.query_moe.aux_loss is not None:
            losses['moe_balance'] = self.query_moe.aux_loss

        # Contrastive learning loss (alignment)
        if self.contrastive_encoder is not None:
            if self.contrastive_encoder.contrastive_loss is not None:
                losses['contrastive'] = (
                    self.contrastive_weight * self.contrastive_encoder.contrastive_loss
                )
            # Uniformity loss (embedding spread on hypersphere)
            if self.contrastive_encoder.uniformity_loss_value is not None:
                losses['uniformity'] = self.contrastive_encoder.uniformity_loss_value

        # VIB KL loss from feature alignment
        if self.use_variational:
            kl_loss = self.model_align.get_kl_loss()
            if kl_loss is not None and kl_loss.item() > 0:
                losses['vib_kl'] = self.vib_beta * kl_loss

        return losses
