import mlx.core as mx
import mlx.nn as nn

from attention import RelPositionMultiHeadAttention

class Swish(nn.Module):
    def __call__(self, x):
        return x * mx.sigmoid(x)

class ConformerFeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0):
        super().__init__()
        self.d_model = d_model
        self.d_ff = d_ff
        
        self.linear1 = nn.Linear(d_model, d_ff)
        self.activation = Swish()
        self.dropout = nn.Dropout(p=dropout)
        self.linear2 = nn.Linear(d_ff, d_model)

    def __call__(self, x):
        x = self.linear1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.linear2(x)
        return x

class ConformerConvolution(nn.Module):
    def __init__(self, d_model: int, kernel_size: int):
        super().__init__()
        self.d_model = d_model
        self.kernel_size = kernel_size
        
        # PyTorch uses (N, C, L) for Conv1d, MLX uses (N, L, C)
        # So pointwise convolutions are just linear layers in MLX
        self.pointwise_conv1 = nn.Linear(d_model, d_model * 2)
        
        # Depthwise Conv1d in MLX
        # groups = channels
        self.depthwise_conv = nn.Conv1d(
            in_channels=d_model,
            out_channels=d_model,
            kernel_size=kernel_size,
            stride=1,
            padding=(kernel_size - 1) // 2,
            groups=d_model
        )
        
        self.batch_norm = nn.BatchNorm(d_model)
        self.activation = Swish()
        self.pointwise_conv2 = nn.Linear(d_model, d_model)

    def __call__(self, x):
        # x is (B, T, C)
        x = self.pointwise_conv1(x)
        
        # GLU activation
        # Split along channels (last dim)
        x1, x2 = mx.split(x, 2, axis=-1)
        x = x1 * mx.sigmoid(x2)
        
        x = self.depthwise_conv(x)
        x = self.batch_norm(x)
        x = self.activation(x)
        x = self.pointwise_conv2(x)
        
        return x

class ConformerLayer(nn.Module):
    def __init__(self, d_model: int, d_ff: int, n_heads: int, conv_kernel_size: int, dropout: float = 0.1):
        super().__init__()
        self.fc_factor = 0.5
        
        # Feed forward 1
        self.norm_feed_forward1 = nn.LayerNorm(d_model)
        self.feed_forward1 = ConformerFeedForward(d_model, d_ff, dropout)
        
        # Self Attention
        self.norm_self_att = nn.LayerNorm(d_model)
        self.self_attn = RelPositionMultiHeadAttention(n_head=n_heads, n_feat=d_model, dropout_rate=dropout)
        
        
        # Convolution
        self.norm_conv = nn.LayerNorm(d_model)
        self.conv = ConformerConvolution(d_model, conv_kernel_size)
        
        # Feed forward 2
        self.norm_feed_forward2 = nn.LayerNorm(d_model)
        self.feed_forward2 = ConformerFeedForward(d_model, d_ff, dropout)
        
        self.dropout = nn.Dropout(dropout)
        self.norm_out = nn.LayerNorm(d_model)

    def __call__(self, x, pos_emb, mask=None):
        # 1. FFN 1
        residual = x
        x = self.norm_feed_forward1(x)
        x = self.feed_forward1(x)
        x = residual + self.dropout(x) * self.fc_factor
        
        # 2. Attention
        residual = x
        x = self.norm_self_att(x)
        x = self.self_attn(query=x, key=x, value=x, mask=mask, pos_emb=pos_emb)
        x = residual + self.dropout(x)
        
        # 3. Convolution
        residual = x
        x = self.norm_conv(x)
        x = self.conv(x)
        x = residual + self.dropout(x)
        
        # 4. FFN 2
        residual = x
        x = self.norm_feed_forward2(x)
        x = self.feed_forward2(x)
        x = residual + self.dropout(x) * self.fc_factor
        
        x = self.norm_out(x)
        return x