"""Neural network building blocks for the GNN router.

Reusable layers and modules used by the GNN architectures. These include
feature alignment, activation functions, gating mechanisms, mixture of experts,
hypernetworks, and contrastive learning components.
"""

from __future__ import annotations

import logging

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

logger = logging.getLogger(__name__)


class SwiGLU(nn.Module):
    """SwiGLU activation function used in LLaMA, PaLM, Mistral.

    More expressive than GELU through learned gating:
        output = (x @ W1) * silu(x @ W_gate) @ W2

    Reference: Shazeer (2020) "GLU Variants Improve Transformer"
    """

    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int) -> None:
        super().__init__()
        self.w1 = nn.Linear(in_dim, hidden_dim, bias=False)
        self.w_gate = nn.Linear(in_dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, out_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w_gate(x)) * self.w1(x))


class ResidualMLP(nn.Module):
    """MLP block with residual connection and pre-normalization.

    Architecture: x + MLP(LayerNorm(x))

    This follows the pre-norm pattern from modern transformers, which provides
    better gradient flow than post-norm for deep networks.

    Args:
        dim: Input/output dimension (must match for residual).
        hidden_mult: Hidden layer multiplier (default 4x like transformers).
        dropout: Dropout rate after activation.
        use_swiglu: Use SwiGLU instead of GELU activation.
    """

    def __init__(
        self,
        dim: int,
        hidden_mult: int = 4,
        dropout: float = 0.1,
        use_swiglu: bool = False,
    ) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        hidden_dim = dim * hidden_mult

        if use_swiglu:
            self.mlp = nn.Sequential(
                SwiGLU(dim, hidden_dim, dim),
                nn.Dropout(dropout),
            )
        else:
            self.mlp = nn.Sequential(
                nn.Linear(dim, hidden_dim),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim, dim),
                nn.Dropout(dropout),
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.mlp(self.norm(x))


class VariationalLayer(nn.Module):
    """Variational Information Bottleneck layer for embeddings.

    Learns a compressed, stochastic representation that retains only
    task-relevant information. During training, samples from q(z|x);
    during inference, uses the mean.

    The KL loss should be added to the task loss with a weight beta:
        total_loss = task_loss + beta * kl_loss

    Reference: Alemi et al. (2017) "Deep Variational Information Bottleneck"

    Args:
        in_dim: Input dimension.
        latent_dim: Latent space dimension (can be same as in_dim).
        min_logvar: Minimum log variance to prevent collapse (default -10).
    """

    def __init__(
        self,
        in_dim: int,
        latent_dim: int,
        min_logvar: float = -10.0,
    ) -> None:
        super().__init__()
        self.min_logvar = min_logvar
        self.fc_mu = nn.Linear(in_dim, latent_dim)
        self.fc_logvar = nn.Linear(in_dim, latent_dim)

        # Store KL loss for retrieval during training
        self.kl_loss: torch.Tensor | None = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mu = self.fc_mu(x)
        logvar = self.fc_logvar(x).clamp(min=self.min_logvar)

        # Reparameterization trick: z = mu + eps * std
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std

            # KL divergence to N(0, I): 0.5 * sum(mu^2 + var - 1 - logvar)
            self.kl_loss = 0.5 * torch.mean(mu.pow(2) + logvar.exp() - 1 - logvar)
        else:
            z = mu
            self.kl_loss = None

        return z


class ModelEMA:
    """Exponential Moving Average of model weights for better generalization.

    Maintains a shadow copy of model weights updated via:
        shadow = decay * shadow + (1 - decay) * model_weights

    Using EMA weights for evaluation often yields 1-2% improvement.
    Based on TensorFlow's ExponentialMovingAverage.

    Args:
        model: The model to track.
        decay: EMA decay rate (0.999 is typical, higher = slower update).
        device: Device to store shadow weights on.
    """

    def __init__(
        self,
        model: nn.Module,
        decay: float = 0.999,
        device: torch.device | None = None,
    ) -> None:
        self.decay = decay
        self.device = device
        # Create shadow copy of parameters
        self.shadow: dict[str, torch.Tensor] = {}
        self.backup: dict[str, torch.Tensor] = {}

        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()
                if device is not None:
                    self.shadow[name] = self.shadow[name].to(device)

    def update(self, model: nn.Module) -> None:
        """Update shadow weights with current model weights."""
        for name, param in model.named_parameters():
            if param.requires_grad and name in self.shadow:
                self.shadow[name].mul_(self.decay).add_(
                    param.data, alpha=1 - self.decay
                )

    def apply_shadow(self, model: nn.Module) -> None:
        """Replace model weights with shadow weights (for evaluation)."""
        for name, param in model.named_parameters():
            if param.requires_grad and name in self.shadow:
                self.backup[name] = param.data.clone()
                param.data.copy_(self.shadow[name])

    def restore(self, model: nn.Module) -> None:
        """Restore original model weights (after evaluation)."""
        for name, param in model.named_parameters():
            if param.requires_grad and name in self.backup:
                param.data.copy_(self.backup[name])
        self.backup.clear()

    def state_dict(self) -> dict[str, torch.Tensor]:
        """Return shadow weights for checkpointing."""
        return self.shadow.copy()

    def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
        """Load shadow weights from checkpoint."""
        self.shadow = state_dict.copy()


class FeatureAlign(nn.Module):
    """Aligns query and LLM features to a common embedding space.

    Supports deep MLP preprocessing for both query and LLM features with optional
    residual connections, SwiGLU activation, and variational information bottleneck
    for improved gradient flow, expressiveness, and regularization.

    Args:
        query_feature_dim: Dimension of input query embeddings.
        llm_feature_dim: Dimension of input LLM embeddings.
        common_dim: Output dimension for aligned features.
        query_mlp_layers: Number of MLP layers for query transform (1 = linear).
        llm_mlp_layers: Number of MLP layers for LLM transform (1 = linear).
        input_dropout: Dropout rate applied to input embeddings (data augmentation).
        use_residual: Use residual MLP blocks instead of plain MLP.
        use_swiglu: Use SwiGLU activation (only with use_residual=True).
        use_variational: Apply variational information bottleneck after transforms.
    """

    def __init__(
        self,
        query_feature_dim: int,
        llm_feature_dim: int,
        common_dim: int,
        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,
    ) -> None:
        super().__init__()
        self.input_dropout = nn.Dropout(input_dropout) if input_dropout > 0 else None
        self.use_variational = use_variational

        # Build query transform
        self.query_transform = self._build_transform(
            query_feature_dim, common_dim, query_mlp_layers, use_residual, use_swiglu
        )

        # Build LLM transform (output is 2x common_dim for bipartite graph)
        self.llm_transform = self._build_transform(
            llm_feature_dim, common_dim * 2, llm_mlp_layers, use_residual, use_swiglu
        )

        self.task_transform = nn.Linear(llm_feature_dim, common_dim)

        # Optional variational bottleneck after transforms
        if use_variational:
            self.query_vib = VariationalLayer(common_dim, common_dim)
            self.llm_vib = VariationalLayer(common_dim * 2, common_dim * 2)
        else:
            self.query_vib = None
            self.llm_vib = None

    def get_kl_loss(self) -> torch.Tensor:
        """Get combined KL loss from variational layers (for training loss).

        Returns a tensor on the same device as the VIB layers. If variational
        is disabled or no KL loss is available, returns zero on the correct device.
        """
        # Determine device from module parameters for consistency
        device = next(self.parameters()).device

        if not self.use_variational:
            return torch.zeros(1, device=device).squeeze()

        kl_loss = torch.zeros(1, device=device).squeeze()
        if self.query_vib is not None and self.query_vib.kl_loss is not None:
            kl_loss = kl_loss + self.query_vib.kl_loss
        if self.llm_vib is not None and self.llm_vib.kl_loss is not None:
            kl_loss = kl_loss + self.llm_vib.kl_loss
        return kl_loss

    def _build_transform(
        self,
        in_dim: int,
        out_dim: int,
        num_layers: int,
        use_residual: bool,
        use_swiglu: bool,
    ) -> nn.Module:
        """Build MLP transform with optional residual connections."""
        if num_layers <= 1:
            return nn.Linear(in_dim, out_dim)

        if use_residual:
            # Residual MLP: project to out_dim, then apply N-1 residual blocks
            # Use hidden_mult=2 (not 4) to avoid OOM on large graphs
            layers: list[nn.Module] = [nn.Linear(in_dim, out_dim)]
            for _ in range(num_layers - 1):
                layers.append(ResidualMLP(out_dim, hidden_mult=2, use_swiglu=use_swiglu))
            return nn.Sequential(*layers)
        else:
            # Original deep MLP (V6 style)
            layers = []
            curr_dim = in_dim
            hidden_dim = (in_dim + out_dim) // 2
            for _ in range(num_layers - 1):
                layers.extend([
                    nn.Linear(curr_dim, hidden_dim),
                    nn.LayerNorm(hidden_dim),
                    nn.GELU(),
                    nn.Dropout(0.1),
                ])
                curr_dim = hidden_dim
            layers.append(nn.Linear(curr_dim, out_dim))
            return nn.Sequential(*layers)

    def forward(
        self,
        task_features: torch.Tensor,
        query_features: torch.Tensor,
        llm_features: torch.Tensor,
    ) -> torch.Tensor:
        # Apply input dropout for data augmentation during training
        if self.input_dropout is not None:
            query_features = self.input_dropout(query_features)

        aligned_task = self.task_transform(task_features)
        aligned_query = self.query_transform(query_features)

        # Apply variational bottleneck if enabled (stochastic during training)
        if self.query_vib is not None:
            aligned_query = self.query_vib(aligned_query)

        query_node_features = torch.cat([aligned_task, aligned_query], dim=1)
        aligned_llm = self.llm_transform(llm_features)

        # Apply variational bottleneck to LLM features
        if self.llm_vib is not None:
            aligned_llm = self.llm_vib(aligned_llm)

        return torch.cat([query_node_features, aligned_llm], dim=0)


class GatedIdentitySkip(nn.Module):
    """Learnable interpolation between raw and GNN-processed features.

    For edge prediction, we currently use x_initial (raw) for queries and x_out
    (GNN-processed) for LLMs. This module allows the model to learn the optimal
    blend: alpha * x_initial + (1-alpha) * x_out.

    This is particularly useful when:
    - Some queries benefit from GNN context (relational)
    - Some queries work better with raw features (intrinsic)

    The gate can be per-feature (dim > 1) or scalar (dim = 1).
    """

    def __init__(self, dim: int, per_feature: bool = True) -> None:
        """Initialize gated skip connection.

        Args:
            dim: Feature dimension for per-feature gating.
            per_feature: If True, learn separate gate per feature. If False,
                use a single scalar gate for all features.
        """
        super().__init__()
        gate_dim = dim if per_feature else 1
        # Initialize to 0.5 (equal mix) via sigmoid of 0
        self.gate_logit = nn.Parameter(torch.zeros(gate_dim))

    def forward(
        self,
        x_initial: torch.Tensor,
        x_out: torch.Tensor,
    ) -> torch.Tensor:
        """Blend raw and processed features with learned gate.

        Args:
            x_initial: Pre-GNN features (num_nodes, dim).
            x_out: Post-GNN features (num_nodes, dim).

        Returns:
            Blended features: alpha * x_initial + (1-alpha) * x_out.
        """
        alpha = torch.sigmoid(self.gate_logit)
        return alpha * x_initial + (1 - alpha) * x_out


class QueryMoE(nn.Module):
    """Mixture of Experts for query-type-specific processing.

    Different query types (code, math, reasoning, factual) may benefit from
    specialized processing. This MoE routes queries through expert sub-networks
    based on a learned router, while sharing the LLM encoder.

    Architecture:
    - Router: Linear(dim, num_experts) -> softmax
    - Experts: num_experts x MLP(dim -> dim)
    - Output: Weighted sum of expert outputs

    Uses top-k routing with load balancing loss for efficient training.
    """

    def __init__(
        self,
        dim: int,
        num_experts: int = 4,
        top_k: int = 2,
        expert_hidden_mult: int = 2,
        dropout: float = 0.1,
    ) -> None:
        """Initialize Mixture of Experts.

        Args:
            dim: Input/output dimension.
            num_experts: Number of expert networks (default 4 for code/math/reasoning/other).
            top_k: Number of experts to activate per query (default 2).
            expert_hidden_mult: Hidden dimension multiplier for experts.
            dropout: Dropout rate within experts.
        """
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Router: predicts expert weights
        self.router = nn.Linear(dim, num_experts)

        # Expert networks: specialized MLPs
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(dim, dim * expert_hidden_mult),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(dim * expert_hidden_mult, dim),
            )
            for _ in range(num_experts)
        ])

        # Load balancing loss for even expert utilization
        self.aux_loss: torch.Tensor | None = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Route queries through top-k experts.

        Args:
            x: Query features (num_queries, dim).

        Returns:
            Expert-processed features (num_queries, dim).
        """
        num_queries, dim = x.shape

        # Compute router logits and select top-k experts
        router_logits = self.router(x)  # (num_queries, num_experts)
        router_probs = F.softmax(router_logits, dim=-1)

        # Top-k selection
        top_k_probs, top_k_indices = router_probs.topk(self.top_k, dim=-1)
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)  # Renormalize

        # Compute load balancing auxiliary loss
        if self.training:
            expert_usage = router_probs.mean(dim=0)
            target_usage = torch.ones_like(expert_usage) / self.num_experts
            self.aux_loss = F.mse_loss(expert_usage, target_usage)
        else:
            self.aux_loss = None

        # Flatten top-k assignments for batched processing
        # Shape: (num_queries * top_k,)
        flat_indices = top_k_indices.flatten()
        flat_weights = top_k_probs.flatten().unsqueeze(1)  # (num_queries * top_k, 1)

        # Replicate input for each top-k slot: (num_queries * top_k, dim)
        x_repeated = x.unsqueeze(1).expand(-1, self.top_k, -1).reshape(-1, dim)

        # Process each expert's batch in a single forward pass (avoids inner loop)
        expert_outputs = torch.zeros_like(x_repeated)
        for e in range(self.num_experts):
            mask = flat_indices == e
            if mask.any():
                expert_outputs[mask] = self.experts[e](x_repeated[mask])

        # Weight and reshape back: (num_queries * top_k, dim) -> (num_queries, top_k, dim)
        weighted_outputs = (flat_weights * expert_outputs).view(num_queries, self.top_k, dim)

        # Sum across top-k experts
        return weighted_outputs.sum(dim=1)


class PreferenceHypernetwork(nn.Module):
    """Hypernetwork that generates routing parameters from preferences.

    Instead of concatenating preference features to edge features, this
    generates the final routing layer's weights/biases conditioned on preferences.
    This provides cleaner preference conditioning without polluting edge features.

    Architecture:
    - Input: preference vector [perf_weight, cost_weight, latency_weight]
    - Output: weight matrix and bias for final routing projection

    Reference: Ha et al. (2017) "Hypernetworks"
    """

    def __init__(
        self,
        pref_dim: int = 3,
        hidden_dim: int = 64,
        output_weight_shape: tuple[int, int] = (256, 1),
    ) -> None:
        """Initialize hypernetwork.

        Args:
            pref_dim: Dimension of preference vector (default 3: perf/cost/latency).
            hidden_dim: Hidden dimension of hypernetwork.
            output_weight_shape: Shape of generated weight matrix (input_dim, output_dim).
        """
        super().__init__()
        self.output_weight_shape = output_weight_shape
        weight_size = output_weight_shape[0] * output_weight_shape[1]
        bias_size = output_weight_shape[1]

        # Hypernetwork: preferences -> routing parameters
        self.weight_generator = nn.Sequential(
            nn.Linear(pref_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, weight_size),
        )
        self.bias_generator = nn.Sequential(
            nn.Linear(pref_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, bias_size),
        )

    def forward(
        self,
        preferences: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Generate routing parameters from preferences.

        Args:
            preferences: Preference vector (pref_dim,) or (batch, pref_dim).

        Returns:
            Tuple of (weight, bias) for routing layer.
            weight: (input_dim, output_dim)
            bias: (output_dim,)
        """
        # Handle both batched and unbatched preferences
        if preferences.dim() == 1:
            preferences = preferences.unsqueeze(0)

        weight = self.weight_generator(preferences)
        bias = self.bias_generator(preferences)

        # Reshape weight to matrix form
        weight = weight.view(*weight.shape[:-1], *self.output_weight_shape)
        return weight.squeeze(0), bias.squeeze(0)


def uniformity_loss(embeddings: torch.Tensor, t: float = 2.0) -> torch.Tensor:
    """Compute uniformity loss encouraging embeddings to spread on hypersphere.

    From Wang & Isola (2020): "Understanding Contrastive Representation Learning
    through Alignment and Uniformity on the Hypersphere"

    Args:
        embeddings: L2-normalized embeddings (batch_size, dim).
        t: Temperature parameter (higher = stronger penalty for clustering).

    Returns:
        Scalar uniformity loss (lower = more uniform distribution).
    """
    sq_pdist = torch.cdist(embeddings, embeddings, p=2).pow(2)
    return sq_pdist.mul(-t).exp().mean().log()


class ContrastiveQueryEncoder(nn.Module):
    """Contrastive learning module for query embeddings.

    Pre-trains query encoder to pull queries toward their best-performing LLM
    and push away from worst-performing LLM. This creates embeddings with
    better structure for routing decisions.

    Loss components:
    - Alignment (InfoNCE): Pull queries toward best LLM, push from worst
    - Uniformity (optional): Spread embeddings uniformly on hypersphere

    This improves the 1024D generic embeddings by making them routing-aware.
    """

    def __init__(
        self,
        query_dim: int,
        llm_dim: int,
        projection_dim: int = 128,
        temperature: float = 0.07,
        use_uniformity: bool = False,
        uniformity_weight: float = 0.05,
        uniformity_t: float = 2.0,
    ) -> None:
        """Initialize contrastive encoder.

        Args:
            query_dim: Query embedding dimension.
            llm_dim: LLM embedding dimension.
            projection_dim: Dimension of contrastive projection space.
            temperature: Temperature for InfoNCE softmax (lower = sharper).
            use_uniformity: Enable uniformity loss to prevent embedding collapse.
            uniformity_weight: Weight for uniformity loss term.
            uniformity_t: Temperature for uniformity kernel (higher = stronger).
        """
        super().__init__()
        self.temperature = temperature
        self.use_uniformity = use_uniformity
        self.uniformity_weight = uniformity_weight
        self.uniformity_t = uniformity_t

        # Project to shared contrastive space
        self.query_proj = nn.Sequential(
            nn.Linear(query_dim, projection_dim),
            nn.ReLU(),
            nn.Linear(projection_dim, projection_dim),
        )
        self.llm_proj = nn.Sequential(
            nn.Linear(llm_dim, projection_dim),
            nn.ReLU(),
            nn.Linear(projection_dim, projection_dim),
        )

        self.contrastive_loss: torch.Tensor | None = None
        self.uniformity_loss_value: torch.Tensor | None = None

    def forward(
        self,
        query_emb: torch.Tensor,
        llm_emb: torch.Tensor,
        best_llm_idx: torch.Tensor | None = None,
        worst_llm_idx: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Project queries and optionally compute contrastive + uniformity loss.

        Args:
            query_emb: Query embeddings (num_queries, query_dim).
            llm_emb: LLM embeddings (num_llms, llm_dim).
            best_llm_idx: Index of best LLM per query for contrastive loss.
            worst_llm_idx: Index of worst LLM per query for hard negatives.

        Returns:
            Projected query embeddings (num_queries, projection_dim).
        """
        # Project to contrastive space
        q_proj = F.normalize(self.query_proj(query_emb), dim=-1)
        l_proj = F.normalize(self.llm_proj(llm_emb), dim=-1)

        self.clear_losses()

        # Compute contrastive loss during training if labels provided
        if self.training and best_llm_idx is not None:
            # Similarity matrix: (num_queries, num_llms)
            sim_matrix = torch.matmul(q_proj, l_proj.T) / self.temperature

            # InfoNCE loss: maximize similarity to best LLM
            self.contrastive_loss = F.cross_entropy(sim_matrix, best_llm_idx)

            # Uniformity loss: spread embeddings on hypersphere
            if self.use_uniformity:
                q_uniform = uniformity_loss(q_proj, self.uniformity_t)
                l_uniform = uniformity_loss(l_proj, self.uniformity_t)
                self.uniformity_loss_value = self.uniformity_weight * (q_uniform + l_uniform)

        return q_proj

    def clear_losses(self) -> None:
        """Clear cached loss tensors from the previous forward pass."""
        self.contrastive_loss = None
        self.uniformity_loss_value = None
