"""GNN Router baseline wrapper.

This module wraps the existing GATEncoderDecoderNet to conform to the BaseRouter
interface, enabling fair comparison with Multi-Head MLP and Two Tower baselines.

Architecture:
    Query nodes + Model nodes -> [Bipartite Graph] -> [GAT Message Passing] -> Edge Scores

    The graph structure:
    - Query nodes: Encoded query embeddings
    - Model nodes: Encoded LLM embeddings
    - Edges: Query-to-Model connections (fully connected bipartite)
    - Message passing: GAT attention aggregates neighborhood information

Key Characteristics:
    - Transductive: Uses full training graph at inference time
    - Message passing enables indirect query-query similarity via shared LLM neighbors
    - Most expressive architecture but highest computational cost

Advantages:
    - Can learn complex relational patterns
    - Attention weights provide interpretability
    - Leverages structure in query-model relationships

Disadvantages:
    - Requires full graph construction at inference
    - Computationally expensive (O(N*M) edges for N queries, M models)
    - Complex training with multiple auxiliary losses
    - Fully-connected bipartite graph may not have meaningful structure

When to use:
    - Research/experimentation with relational learning
    - When you have meaningful graph structure to exploit
    - When training data is sparse (transfer via graph structure)
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data

from .base import (
    BaseRouter,
    DifficultySignals,
    ResidualBlock,
    RouterConfig,
    RouterOutput,
    SEBlock,
)


@dataclass
class GNNRouterConfig(RouterConfig):
    """Extended configuration for GNN router.

    Inherits from RouterConfig and adds GNN-specific parameters.

    Attributes:
        num_gnn_layers: Number of GAT message passing layers.
        attention_dropout: Dropout for attention coefficients.
        edge_feature_dim: Dimension of edge features (preferences, etc.).
        use_gated_skip: Learnable blend of pre/post-GNN features.
        use_contrastive: Enable contrastive learning auxiliary loss.
        contrastive_weight: Weight for contrastive loss term.
    """

    num_gnn_layers: int = 3
    attention_dropout: float = 0.1
    edge_feature_dim: int = 6
    use_gated_skip: bool = False
    use_contrastive: bool = True
    contrastive_weight: float = 0.1


class SimplifiedGNNRouter(BaseRouter):
    """Simplified GNN router for baseline comparison.

    This is a clean reimplementation of the core GNN routing logic,
    without the complexity of the full GATEncoderDecoderNet.

    The simplification:
    1. No training graph retention (fresh graph per forward pass)
    2. No auxiliary losses (contrastive, uniformity, etc.)
    3. Standard GAT message passing without advanced features
    4. Clean interface matching other baselines
    """

    def __init__(self, config: RouterConfig) -> None:
        """Initialize simplified GNN router.

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

        # Check for PyTorch Geometric
        try:
            from torch_geometric.nn import GATv2Conv
        except ImportError:
            raise ImportError(
                "torch_geometric is required for GNN router. "
                "Install with: pip install torch_geometric"
            )

        # Query encoder (projects to hidden space)
        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 to hidden space)
        self.model_encoder = nn.Sequential(
            nn.Linear(config.model_dim, config.hidden_dim),
            nn.LayerNorm(config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            ResidualBlock(config.hidden_dim, dropout=config.dropout),
        )

        # Optional SE blocks
        if config.use_se_block:
            self.query_se = SEBlock(config.hidden_dim, reduction=4)
            self.model_se = SEBlock(config.hidden_dim, reduction=4)
        else:
            self.query_se = nn.Identity()
            self.model_se = nn.Identity()

        # GAT layers for message passing
        self.gat_layers = nn.ModuleList()
        self.gat_norms = nn.ModuleList()

        for _ in range(config.num_layers):
            self.gat_layers.append(
                GATv2Conv(
                    in_channels=config.hidden_dim,
                    out_channels=config.hidden_dim,
                    heads=config.num_heads,
                    concat=False,  # Average heads
                    dropout=config.dropout,
                    add_self_loops=False,  # Bipartite graph
                )
            )
            self.gat_norms.append(nn.LayerNorm(config.hidden_dim))

        # Edge scoring MLP
        self.edge_scorer = nn.Sequential(
            nn.Linear(config.hidden_dim * 2, config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.hidden_dim, 1),
        )

        # Temperature
        self.log_temperature = nn.Parameter(
            torch.log(torch.tensor(config.temperature))
        )

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

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

    def _build_bipartite_graph(
        self,
        num_queries: int,
        num_models: int,
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build fully-connected bipartite edge index.

        Args:
            num_queries: Number of query nodes.
            num_models: Number of model nodes.
            device: Device to create tensors on.

        Returns:
            Tuple of (edge_index, query_indices, model_indices) where:
            - edge_index: (2, num_queries * num_models) edge connections
            - query_indices: Query node index for each edge
            - model_indices: Model node index for each edge
        """
        # Query nodes: 0 to num_queries-1
        # Model nodes: num_queries to num_queries+num_models-1

        query_indices = torch.arange(num_queries, device=device)
        model_indices = torch.arange(num_models, device=device) + num_queries

        # Create all query->model edges
        src = query_indices.repeat_interleave(num_models)
        dst = model_indices.repeat(num_queries)

        edge_index = torch.stack([src, dst], dim=0)

        return edge_index, query_indices, model_indices

    def compute_scores(
        self, query_repr: torch.Tensor, model_repr: torch.Tensor
    ) -> torch.Tensor:
        """Compute scores via GNN message passing.

        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]
        device = query_repr.device

        # Concatenate query and model nodes
        # Node ordering: queries first, then models
        x = torch.cat([query_repr, model_repr], dim=0)

        # Build bipartite graph edges
        edge_index, _, _ = self._build_bipartite_graph(
            batch_size, num_models, device
        )

        # Apply GAT layers with residual connections
        for gat, norm in zip(self.gat_layers, self.gat_norms):
            x_new = gat(x, edge_index)
            x_new = norm(x_new)
            x = x + F.gelu(x_new)  # Residual + activation

        # Extract updated representations
        query_out = x[:batch_size]  # (batch, hidden)
        model_out = x[batch_size:]  # (num_models, hidden)

        # Compute edge scores via MLP on concatenated features
        # For each (query, model) pair
        query_expanded = query_out.unsqueeze(1).expand(-1, num_models, -1)
        model_expanded = model_out.unsqueeze(0).expand(batch_size, -1, -1)
        edge_features = torch.cat([query_expanded, model_expanded], dim=-1)

        scores = self.edge_scorer(edge_features).squeeze(-1)  # (batch, num_models)

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

        return scores


class GNNRouter(BaseRouter):
    """GNN router wrapper for the production GATEncoderDecoderNet.

    This wrapper adapts the full GATEncoderDecoderNet to the BaseRouter interface.
    It preserves all production features while enabling comparison with simpler baselines.

    Note: This wrapper requires the full training graph to be loaded for inference,
    matching the production behavior.
    """

    def __init__(
        self,
        config: RouterConfig,
        gnn_model: Optional[nn.Module] = None,
    ) -> None:
        """Initialize GNN router wrapper.

        Args:
            config: RouterConfig with architecture parameters.
            gnn_model: Pre-initialized GATEncoderDecoderNet (optional).
                      If None, creates a new SimplifiedGNNRouter.
        """
        super().__init__(config)

        if gnn_model is not None:
            self.gnn = gnn_model
            self._use_production_gnn = True
        else:
            # Fall back to simplified GNN for baseline comparison
            self.gnn = SimplifiedGNNRouter(config)
            self._use_production_gnn = False

        # Cache for graph data
        self._cached_graph: Optional[Data] = None

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

        For production GNN, this returns raw embeddings (encoding happens in forward).
        For simplified GNN, delegates to the internal encoder.
        """
        if self._use_production_gnn:
            return query_emb
        else:
            return self.gnn.encode_query(query_emb)

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

        For production GNN, this returns raw embeddings (encoding happens in forward).
        For simplified GNN, delegates to the internal encoder.
        """
        if self._use_production_gnn:
            return model_embs
        else:
            return self.gnn.encode_models(model_embs)

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

        Args:
            query_repr: Encoded queries (or raw for production GNN).
            model_repr: Encoded models (or raw for production GNN).

        Returns:
            Routing scores, shape (batch, num_models).
        """
        if self._use_production_gnn:
            raise NotImplementedError(
                "Production GNN requires full graph context. "
                "Use forward_with_graph() instead."
            )
        else:
            return self.gnn.compute_scores(query_repr, model_repr)

    def forward(
        self,
        query_emb: torch.Tensor,
        model_embs: torch.Tensor,
        return_signals: bool = False,
    ) -> RouterOutput:
        """Forward pass for simplified GNN.

        For production GNN, use forward_with_graph() instead.
        """
        if self._use_production_gnn:
            raise NotImplementedError(
                "Production GNN requires full graph context. "
                "Use forward_with_graph() instead."
            )

        return self.gnn.forward(query_emb, model_embs, return_signals)

    def forward_with_graph(
        self,
        data: Data,
        edge_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, Optional[DifficultySignals]]:
        """Forward pass with full graph context (production GNN).

        Args:
            data: PyG Data object with full graph.
            edge_mask: Mask for edges to predict.

        Returns:
            Tuple of (scores, signals) for masked edges.
        """
        if not self._use_production_gnn:
            raise ValueError(
                "forward_with_graph() requires production GNN. "
                "Use forward() for simplified baseline."
            )

        # This would call the production GNN's forward method
        # Implementation depends on the specific GATEncoderDecoderNet API
        raise NotImplementedError(
            "Production GNN integration requires additional setup. "
            "Use SimplifiedGNNRouter for baseline comparison."
        )

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_path: str,
        config: RouterConfig,
    ) -> "GNNRouter":
        """Load GNN router from checkpoint.

        Args:
            checkpoint_path: Path to model checkpoint.
            config: RouterConfig for initialization.

        Returns:
            GNNRouter with loaded weights.
        """
        from embergraphrouter.model.graph_nn import GATEncoderDecoderNet

        # Load checkpoint
        checkpoint = torch.load(checkpoint_path, map_location="cpu")

        # Extract model config from checkpoint
        model_config = checkpoint.get("config", {})

        # Initialize production GNN
        gnn = GATEncoderDecoderNet(
            query_feature_dim=model_config.get("query_dim", config.query_dim),
            llm_feature_dim=model_config.get("model_dim", config.model_dim),
            hidden_features=model_config.get("hidden_dim", config.hidden_dim),
            in_edges=model_config.get("edge_feature_dim", 6),
            num_gnn_layers=model_config.get("num_gnn_layers", config.num_layers),
            num_attention_heads=model_config.get("num_heads", config.num_heads),
        )

        # Load weights
        gnn.load_state_dict(checkpoint["model_state_dict"])

        return cls(config, gnn_model=gnn)


# Utility function for creating GNN data objects
def create_routing_graph(
    query_embs: torch.Tensor,
    model_embs: torch.Tensor,
    targets: Optional[torch.Tensor] = None,
    edge_features: Optional[torch.Tensor] = None,
) -> Data:
    """Create PyG Data object for GNN routing.

    Args:
        query_embs: Query embeddings, shape (num_queries, query_dim).
        model_embs: Model embeddings, shape (num_models, model_dim).
        targets: Optional target scores, shape (num_queries, num_models).
        edge_features: Optional edge features, shape (num_edges, edge_dim).

    Returns:
        PyG Data object with bipartite graph structure.
    """
    num_queries = query_embs.shape[0]
    num_models = model_embs.shape[0]
    device = query_embs.device

    # Node features: queries then models
    x = torch.cat([query_embs, model_embs], dim=0)

    # Build edge index (fully connected bipartite)
    query_indices = torch.arange(num_queries, device=device)
    model_indices = torch.arange(num_models, device=device) + num_queries

    src = query_indices.repeat_interleave(num_models)
    dst = model_indices.repeat(num_queries)
    edge_index = torch.stack([src, dst], dim=0)

    # Create Data object
    data = Data(x=x, edge_index=edge_index)

    # Add edge features if provided
    if edge_features is not None:
        data.edge_attr = edge_features

    # Add targets if provided
    if targets is not None:
        # Flatten targets to match edge ordering
        data.y = targets.flatten()

    # Store metadata
    data.num_queries = num_queries
    data.num_models = num_models

    return data
