"""
CyMeta Feed-Forward Network Layers.

Drop-in replacements for FFN layers that use CyMeta compression.
"""

import torch
import torch.nn as nn
from typing import Optional
from cymeta.models.linear import CyMetaLinear


class CyMetaFFN(nn.Module):
    """
    Feed-forward network with CyMeta-compressed layers.
    
    This replaces standard FFN up and down projections with CyMeta
    compressed versions. Typically used in transformer blocks.
    
    Args:
        embed_dim: Embedding dimension
        ffn_dim: Feed-forward hidden dimension (typically 4 * embed_dim)
        dict_size: Number of dictionary atoms (default: 32)
        shift_bits: Number of bits for shift encoding (default: 8)
        gating_hidden: Hidden dimension for gating network (default: 16)
        activation: Activation function ('gelu', 'relu', 'swish') (default: 'gelu')
        dropout: Dropout probability (default: 0.0)
        bias: Whether to include bias in projections (default: True)
    """
    
    def __init__(
        self,
        embed_dim: int,
        ffn_dim: int,
        dict_size: int = 32,
        shift_bits: int = 8,
        gating_hidden: int = 16,
        activation: str = "gelu",
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.ffn_dim = ffn_dim
        self.dict_size = dict_size
        self.shift_bits = shift_bits
        self.gating_hidden = gating_hidden
        
        # Up projection (embed_dim -> ffn_dim)
        self.up_proj = CyMetaLinear(
            embed_dim, ffn_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        
        # Down projection (ffn_dim -> embed_dim)
        self.down_proj = CyMetaLinear(
            ffn_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        
        # Activation function
        if activation.lower() == "gelu":
            self.activation = nn.GELU()
        elif activation.lower() == "relu":
            self.activation = nn.ReLU()
        elif activation.lower() == "swish":
            self.activation = nn.SiLU()  # SiLU is Swish
        else:
            raise ValueError(f"Unsupported activation: {activation}")
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(
        self,
        x: torch.Tensor,
        token_embedding: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Forward pass of feed-forward network.
        
        Args:
            x: Input tensor of shape (batch_size, seq_len, embed_dim)
            token_embedding: Optional token embedding for context-dependent gating
        
        Returns:
            Output tensor of shape (batch_size, seq_len, embed_dim)
        """
        # Use input as token embedding if not provided
        if token_embedding is None:
            token_embedding = x
        
        # Up projection
        x = self.up_proj(x, token_embedding=token_embedding)
        
        # Activation
        x = self.activation(x)
        
        # Dropout
        x = self.dropout(x)
        
        # Down projection
        # For down_proj, use the current x (after up_proj) as token embedding
        # since it matches the input dimension of down_proj
        x = self.down_proj(x, token_embedding=x)
        
        return x
    
    def get_compression_ratio(self) -> float:
        """
        Calculate average compression ratio across both projections.
        
        Returns:
            Average compression ratio
        """
        up_ratio = self.up_proj.get_compression_ratio()
        down_ratio = self.down_proj.get_compression_ratio()
        return (up_ratio + down_ratio) / 2.0
    
    def __repr__(self) -> str:
        return (
            f"CyMetaFFN(embed_dim={self.embed_dim}, ffn_dim={self.ffn_dim}, "
            f"dict_size={self.dict_size}, "
            f"compression={self.get_compression_ratio():.2f}x)"
        )

