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

def calc_length(lengths, all_paddings, kernel_size, stride, repeat_num=1):
    """Calculates the output length of a Tensor passed through a convolution layer"""
    add_pad = all_paddings - kernel_size
    for i in range(repeat_num):
        lengths = mx.floor((lengths + add_pad) / stride) + 1
    return lengths.astype(mx.int32)

class ConvSubsampling(nn.Module):
    """
    MLX Translation of NVIDIA NeMo ConvSubsampling
    Used for the front-end of FastConformer.
    """
    def __init__(
        self,
        subsampling: str,
        subsampling_factor: int,
        feat_in: int,
        feat_out: int,
        conv_channels: int,
    ):
        super().__init__()
        self._subsampling = subsampling
        self._conv_channels = conv_channels
        self._feat_in = feat_in
        self._feat_out = feat_out
        
        self._sampling_num = int(math.log(subsampling_factor, 2))
        self.subsampling_factor = subsampling_factor
        
        in_channels = 1
        self.conv = []
        
        if subsampling == 'striding':
            self._stride = 2
            self._kernel_size = 3
            self._left_padding = (self._kernel_size - 1) // 2
            self._right_padding = (self._kernel_size - 1) // 2
            
            for i in range(self._sampling_num):
                self.conv.append(
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=conv_channels,
                        kernel_size=self._kernel_size,
                        stride=self._stride,
                        padding=self._left_padding,
                    )
                )
                in_channels = conv_channels
                
        elif subsampling == 'dw_striding':
            self._stride = 2
            self._kernel_size = 3
            self._left_padding = (self._kernel_size - 1) // 2
            self._right_padding = (self._kernel_size - 1) // 2
            
            # Layer 0: Full Conv
            self.conv.append(
                nn.Conv2d(
                    in_channels=in_channels,
                    out_channels=conv_channels,
                    kernel_size=self._kernel_size,
                    stride=self._stride,
                    padding=self._left_padding,
                )
            )
            # Layer 1: ReLU (not a module with weights, but we need to keep indexing aligned, so we use None or handle it differently)
            self.conv.append(None) # Index 1
            
            in_channels = conv_channels
            
            for i in range(self._sampling_num - 1):
                # Depthwise Conv (Layer 2, 5...)
                self.conv.append(
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=in_channels,
                        kernel_size=self._kernel_size,
                        stride=self._stride,
                        padding=self._left_padding,
                        groups=in_channels,
                    )
                )
                
                # Pointwise Conv (Layer 3, 6...)
                self.conv.append(
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=conv_channels,
                        kernel_size=1,
                        stride=1,
                        padding=0,
                        groups=1,
                    )
                )
                
                # ReLU placeholder (Layer 4, 7...)
                self.conv.append(None)
                in_channels = conv_channels
                
        else:
            raise ValueError(f"Subsampling {subsampling} not implemented in MLX yet")

        # In PyTorch: b, c, t, f -> b, t, c*f
        # In MLX: inputs to Conv2d are expected to be (N, H, W, C)
        # So we map: N=batch, H=time, W=freq, C=channels
        # After convs: flatten the last two dimensions to project to feat_out
        
        # Calculate output frequency size mathematically or dynamically
        dummy_f = mx.array([feat_in], dtype=mx.float32)
        out_f = calc_length(
            dummy_f,
            all_paddings=self._left_padding + self._right_padding,
            kernel_size=self._kernel_size,
            stride=self._stride,
            repeat_num=self._sampling_num
        )
        out_f_val = out_f.item()
        
        self.out = nn.Linear(conv_channels * int(out_f_val), feat_out)

    def __call__(self, x, lengths):
        """
        x: shape (batch, time, freq)
        lengths: shape (batch,)
        """
        # Convert to MLX Conv2d shape: (N, H, W, C) -> (batch, time, freq, 1)
        x = mx.expand_dims(x, axis=-1)
        
        for conv in self.conv:
            if conv is None:
                x = mx.maximum(x, 0.0) # ReLU
            else:
                x = conv(x)
            
        b, t, f, c = x.shape
        
        # In PyTorch, shape is (b, c, t, f) which is transposed to (b, t, c, f) and then flattened to (b, t, c*f).
        # In MLX, shape is (b, t, f, c). We MUST transpose to (b, t, c, f) before flattening to match the trained Linear weights.
        x = mx.transpose(x, (0, 1, 3, 2)) # shape becomes (b, t, c, f)
        
        # Flatten c and f
        x = mx.reshape(x, (b, t, c * f))
        
        x = self.out(x)
        
        out_lengths = calc_length(
            lengths,
            all_paddings=self._left_padding + self._right_padding,
            kernel_size=self._kernel_size,
            stride=self._stride,
            repeat_num=self._sampling_num
        )
        
        return x, out_lengths