"""
Gating Network Module for CyMeta Compression.

This module implements a tiny contextual gating network that outputs mixture
coefficients over dictionary atoms. The gating network enables context-dependent
weight reconstruction, ensuring expressivity even with extreme compression.
"""

import torch
import torch.nn as nn
from typing import Optional, Tuple


class GatingNetwork(nn.Module):
    """
    Tiny contextual gating network for adaptive dictionary atom mixing.
    
    For each token and attention head, this network outputs mixture coefficients
    over dictionary atoms. This produces context-dependent reconstructed weight
    vectors at runtime, ensuring expressivity even with extreme compression.
    
    The network is designed to be extremely lightweight to maintain the
    compression benefits while enabling adaptive reconstruction.
    
    Args:
        input_dim: Dimension of input token embeddings
        num_atoms: Number of dictionary atoms to mix over
        hidden_dim: Hidden dimension of the gating MLP (default: 16)
        num_heads: Number of attention heads (for head-specific gating, optional)
        use_head_specific: Whether to use separate gating per attention head
        activation: Activation function (default: GELU)
    
    Attributes:
        mlp: Multi-layer perceptron for computing gating coefficients
    """
    
    def __init__(
        self,
        input_dim: int,
        num_atoms: int,
        hidden_dim: int = 16,
        num_heads: Optional[int] = None,
        use_head_specific: bool = False,
        activation: str = "gelu",
    ):
        super().__init__()
        self.input_dim = input_dim
        self.num_atoms = num_atoms
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.use_head_specific = use_head_specific
        
        # Select activation function
        if activation.lower() == "gelu":
            self.activation = nn.GELU()
        elif activation.lower() == "relu":
            self.activation = nn.ReLU()
        elif activation.lower() == "tanh":
            self.activation = nn.Tanh()
        else:
            raise ValueError(f"Unsupported activation: {activation}")
        
        # Build lightweight MLP
        if use_head_specific and num_heads is not None:
            # Head-specific gating: separate network per head
            self.mlp = nn.ModuleList([
                self._build_mlp(input_dim, num_atoms, hidden_dim)
                for _ in range(num_heads)
            ])
        else:
            # Shared gating network
            self.mlp = self._build_mlp(input_dim, num_atoms, hidden_dim)
    
    def _build_mlp(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dim: int
    ) -> nn.Module:
        """
        Build a lightweight MLP for gating.
        
        Args:
            input_dim: Input dimension
            output_dim: Output dimension (number of atoms)
            hidden_dim: Hidden dimension
        
        Returns:
            MLP module
        """
        return nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            self.activation,
            nn.Linear(hidden_dim, output_dim),
            nn.Softmax(dim=-1)  # Normalize to mixture coefficients
        )
    
    def forward(
        self,
        token_embedding: torch.Tensor,
        head: Optional[int] = None,
        position: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Compute gating coefficients for dictionary atom mixing.
        
        Args:
            token_embedding: Token embedding tensor of shape (batch_size, seq_len, input_dim)
                            or (batch_size, input_dim) for single token
            head: Attention head index (if use_head_specific=True)
            position: Optional position encoding tensor
        
        Returns:
            Gating coefficients tensor of shape (..., num_atoms)
        """
        # Handle different input shapes
        if token_embedding.dim() == 2:
            # Single token: (batch_size, input_dim)
            x = token_embedding
        elif token_embedding.dim() == 3:
            # Sequence: (batch_size, seq_len, input_dim)
            # Use mean pooling or last token (can be made configurable)
            x = token_embedding.mean(dim=1)  # Mean pooling over sequence
        else:
            raise ValueError(
                f"Expected token_embedding to be 2D or 3D, got {token_embedding.dim()}D"
            )
        
        # Add position encoding if provided
        if position is not None:
            if position.dim() == 1:
                position = position.unsqueeze(0).expand_as(x)
            x = x + position
        
        # Apply head-specific or shared gating
        if self.use_head_specific and self.num_heads is not None:
            if head is None:
                raise ValueError("head must be provided when use_head_specific=True")
            if head < 0 or head >= self.num_heads:
                raise ValueError(f"head must be in [0, {self.num_heads}), got {head}")
            gating_coeffs = self.mlp[head](x)
        else:
            gating_coeffs = self.mlp(x)
        
        return gating_coeffs
    
    def get_num_parameters(self) -> int:
        """
        Get the number of parameters in the gating network.
        
        Returns:
            Total number of parameters
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad)
    
    def __repr__(self) -> str:
        return (
            f"GatingNetwork(input_dim={self.input_dim}, "
            f"num_atoms={self.num_atoms}, hidden_dim={self.hidden_dim}, "
            f"num_heads={self.num_heads}, use_head_specific={self.use_head_specific}, "
            f"params={self.get_num_parameters()})"
        )


class ContextualGating(nn.Module):
    """
    Enhanced contextual gating with additional context features.
    
    This variant can incorporate additional context information such as:
    - Layer index
    - Module type
    - Global attention patterns
    
    Args:
        input_dim: Dimension of input token embeddings
        num_atoms: Number of dictionary atoms
        hidden_dim: Hidden dimension
        context_dim: Dimension of additional context features (default: 0)
    """
    
    def __init__(
        self,
        input_dim: int,
        num_atoms: int,
        hidden_dim: int = 16,
        context_dim: int = 0,
    ):
        super().__init__()
        self.input_dim = input_dim
        self.num_atoms = num_atoms
        self.hidden_dim = hidden_dim
        self.context_dim = context_dim
        
        # Combine token embedding and context
        combined_dim = input_dim + context_dim
        
        self.mlp = nn.Sequential(
            nn.Linear(combined_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, num_atoms),
            nn.Softmax(dim=-1)
        )
    
    def forward(
        self,
        token_embedding: torch.Tensor,
        context: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Compute gating coefficients with optional context.
        
        Args:
            token_embedding: Token embedding tensor
            context: Optional context features tensor
        
        Returns:
            Gating coefficients tensor
        """
        if context is not None:
            if token_embedding.shape[:-1] != context.shape[:-1]:
                raise ValueError(
                    f"Token embedding shape {token_embedding.shape} and "
                    f"context shape {context.shape} must match except last dimension"
                )
            x = torch.cat([token_embedding, context], dim=-1)
        else:
            if self.context_dim > 0:
                # Use zero context if not provided
                batch_shape = token_embedding.shape[:-1]
                zero_context = torch.zeros(
                    (*batch_shape, self.context_dim),
                    device=token_embedding.device,
                    dtype=token_embedding.dtype
                )
                x = torch.cat([token_embedding, zero_context], dim=-1)
            else:
                x = token_embedding
        
        return self.mlp(x)

