from typing import List, Tuple

import torch
import torch.nn as nn

from models.scnet_unofficial.utils import create_intervals


class Downsample(nn.Module):
    """
    Downsample class implements a module for downsampling input tensors using 2D convolution.

    Args:
    - input_dim (int): Dimensionality of the input channels.
    - output_dim (int): Dimensionality of the output channels.
    - stride (int): Stride value for the convolution operation.

    Shapes:
    - Input: (B, C_in, F, T) where
        B is batch size,
        C_in is the number of input channels,
        F is the frequency dimension,
        T is the time dimension.
    - Output: (B, C_out, F // stride, T) where
        B is batch size,
        C_out is the number of output channels,
        F // stride is the downsampled frequency dimension.

    """

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        stride: int,
    ):
        """
        Initializes Downsample with input dimension, output dimension, and stride.
        """
        super().__init__()
        self.conv = nn.Conv2d(input_dim, output_dim, 1, (stride, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs forward pass through the Downsample module.

        Args:
        - x (torch.Tensor): Input tensor of shape (B, C_in, F, T).

        Returns:
        - torch.Tensor: Downsampled tensor of shape (B, C_out, F // stride, T).
        """
        return self.conv(x)


class ConvolutionModule(nn.Module):
    """
    ConvolutionModule class implements a module with a sequence of convolutional layers similar to Conformer.

    Args:
    - input_dim (int): Dimensionality of the input features.
    - hidden_dim (int): Dimensionality of the hidden features.
    - kernel_sizes (List[int]): List of kernel sizes for the convolutional layers.
    - bias (bool, optional): If True, adds a learnable bias to the output. Default is False.

    Shapes:
    - Input: (B, T, D) where
        B is batch size,
        T is sequence length,
        D is input dimensionality.
    - Output: (B, T, D) where
        B is batch size,
        T is sequence length,
        D is input dimensionality.
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        kernel_sizes: List[int],
        bias: bool = False,
    ) -> None:
        """
        Initializes ConvolutionModule with input dimension, hidden dimension, kernel sizes, and bias.
        """
        super().__init__()
        self.sequential = nn.Sequential(
            nn.GroupNorm(num_groups=1, num_channels=input_dim),
            nn.Conv1d(
                input_dim,
                2 * hidden_dim,
                kernel_sizes[0],
                stride=1,
                padding=(kernel_sizes[0] - 1) // 2,
                bias=bias,
            ),
            nn.GLU(dim=1),
            nn.Conv1d(
                hidden_dim,
                hidden_dim,
                kernel_sizes[1],
                stride=1,
                padding=(kernel_sizes[1] - 1) // 2,
                groups=hidden_dim,
                bias=bias,
            ),
            nn.GroupNorm(num_groups=1, num_channels=hidden_dim),
            nn.SiLU(),
            nn.Conv1d(
                hidden_dim,
                input_dim,
                kernel_sizes[2],
                stride=1,
                padding=(kernel_sizes[2] - 1) // 2,
                bias=bias,
            ),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs forward pass through the ConvolutionModule.

        Args:
        - x (torch.Tensor): Input tensor of shape (B, T, D).

        Returns:
        - torch.Tensor: Output tensor of shape (B, T, D).
        """
        x = x.transpose(1, 2)
        x = x + self.sequential(x)
        x = x.transpose(1, 2)
        return x


class SDLayer(nn.Module):
    """
    SDLayer class implements a subband decomposition layer with downsampling and convolutional modules.

    Args:
    - subband_interval (Tuple[float, float]): Tuple representing the frequency interval for subband decomposition.
    - input_dim (int): Dimensionality of the input channels.
    - output_dim (int): Dimensionality of the output channels after downsampling.
    - downsample_stride (int): Stride value for the downsampling operation.
    - n_conv_modules (int): Number of convolutional modules.
    - kernel_sizes (List[int]): List of kernel sizes for the convolutional layers.
    - bias (bool, optional): If True, adds a learnable bias to the convolutional layers. Default is True.

    Shapes:
    - Input: (B, Fi, T, Ci) where
        B is batch size,
        Fi is the number of input subbands,
        T is sequence length, and
        Ci is the number of input channels.
    - Output: (B, Fi+1, T, Ci+1) where
        B is batch size,
        Fi+1 is the number of output subbands,
        T is sequence length,
        Ci+1 is the number of output channels.
    """

    def __init__(
        self,
        subband_interval: Tuple[float, float],
        input_dim: int,
        output_dim: int,
        downsample_stride: int,
        n_conv_modules: int,
        kernel_sizes: List[int],
        bias: bool = True,
    ):
        """
        Initializes SDLayer with subband interval, input dimension,
        output dimension, downsample stride, number of convolutional modules, kernel sizes, and bias.
        """
        super().__init__()
        self.subband_interval = subband_interval
        self.downsample = Downsample(input_dim, output_dim, downsample_stride)
        self.activation = nn.GELU()
        conv_modules = [
            ConvolutionModule(
                input_dim=output_dim,
                hidden_dim=output_dim // 4,
                kernel_sizes=kernel_sizes,
                bias=bias,
            )
            for _ in range(n_conv_modules)
        ]
        self.conv_modules = nn.Sequential(*conv_modules)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs forward pass through the SDLayer.

        Args:
        - x (torch.Tensor): Input tensor of shape (B, Fi, T, Ci).

        Returns:
        - torch.Tensor: Output tensor of shape (B, Fi+1, T, Ci+1).
        """
        B, F, T, C = x.shape
        x = x[:, int(self.subband_interval[0] * F) : int(self.subband_interval[1] * F)]
        x = x.permute(0, 3, 1, 2)
        x = self.downsample(x)
        x = self.activation(x)
        x = x.permute(0, 2, 3, 1)

        B, F, T, C = x.shape
        x = x.reshape((B * F), T, C)
        x = self.conv_modules(x)
        x = x.reshape(B, F, T, C)

        return x


class SDBlock(nn.Module):
    """
    SDBlock class implements a block with subband decomposition layers and global convolution.

    Args:
    - input_dim (int): Dimensionality of the input channels.
    - output_dim (int): Dimensionality of the output channels.
    - bandsplit_ratios (List[float]): List of ratios for splitting the frequency bands.
    - downsample_strides (List[int]): List of stride values for downsampling in each subband layer.
    - n_conv_modules (List[int]): List specifying the number of convolutional modules in each subband layer.
    - kernel_sizes (List[int], optional): List of kernel sizes for the convolutional layers. Default is None.

    Shapes:
    - Input: (B, Fi, T, Ci) where
        B is batch size,
        Fi is the number of input subbands,
        T is sequence length,
        Ci is the number of input channels.
    - Output: (B, Fi+1, T, Ci+1) where
        B is batch size,
        Fi+1 is the number of output subbands,
        T is sequence length,
        Ci+1 is the number of output channels.
    """

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        bandsplit_ratios: List[float],
        downsample_strides: List[int],
        n_conv_modules: List[int],
        kernel_sizes: List[int] = None,
    ):
        """
        Initializes SDBlock with input dimension, output dimension, band split ratios, downsample strides, number of convolutional modules, and kernel sizes.
        """
        super().__init__()
        if kernel_sizes is None:
            kernel_sizes = [3, 3, 1]
        assert sum(bandsplit_ratios) == 1, "The split ratios must sum up to 1."
        subband_intervals = create_intervals(bandsplit_ratios)
        self.sd_layers = nn.ModuleList(
            SDLayer(
                input_dim=input_dim,
                output_dim=output_dim,
                subband_interval=sbi,
                downsample_stride=dss,
                n_conv_modules=ncm,
                kernel_sizes=kernel_sizes,
            )
            for sbi, dss, ncm in zip(
                subband_intervals, downsample_strides, n_conv_modules
            )
        )
        self.global_conv2d = nn.Conv2d(output_dim, output_dim, 1, 1)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Performs forward pass through the SDBlock.

        Args:
        - x (torch.Tensor): Input tensor of shape (B, Fi, T, Ci).

        Returns:
        - Tuple[torch.Tensor, torch.Tensor]: Output tensor and skip connection tensor.
        """
        x_skip = torch.concat([layer(x) for layer in self.sd_layers], dim=1)
        x = self.global_conv2d(x_skip.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
        return x, x_skip