import typing
import torch
from torch import nn

from . import base


class DepthwiseSeperableConv(base.BaseModule):
    """深度可分离卷积层

    将标准卷积分解为深度卷积（Depthwise）和逐点卷积（Pointwise）两个步骤，
    大幅减少参数量和计算量，同时保持相似的表达能力。

    深度卷积：每个输入通道独立进行卷积
    逐点卷积：使用 1x1 卷积进行通道间的线性组合

    Examples
    ----------
    ```py
    # 1D EEG 信号处理
    ds_conv = DepthwiseSeperableConv(
        in_channels=64, out_channels=128, kernel_size=3,
        stride=1, padding=1, Conv=nn.Conv1d
    )
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    output = ds_conv(eeg_data)  # [N, 128, 1000]

    # 2D 图像处理
    ds_conv_2d = DepthwiseSeperableConv(
        in_channels=256, out_channels=512, kernel_size=3,
        stride=2, padding=1, Conv=nn.Conv2d
    )
    image = torch.randn(16, 256, 32, 32)  # [N, C, H, W]
    features = ds_conv_2d(image)  # [N, 512, 16, 16]
    ```
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        bias: bool = False,
        Conv=nn.Conv1d,
    ):
        """初始化 DepthwiseSeperableConv 层

        Parameters
        ----------
        - `in_channels` : `int`
            - 输入通道数
        - `out_channels` : `int`
            - 输出通道数
        - `kernel_size` : `int`
            - 卷积核大小
        - `stride` : `int`, optional
            - 步长, by default 1
        - `padding` : `int`, optional
            - 填充, by default 0
        - `dilation` : `int`, optional
            - 膨胀率, by default 1
        - `bias` : `bool`, optional
            - 是否使用偏置, by default False
        - `Conv` : `nn.Module`, optional
            - 卷积层类型 (nn.Conv1d, nn.Conv2d, nn.Conv3d), by default nn.Conv1d
        """
        super(DepthwiseSeperableConv, self).__init__()
        self.depthwise_conv = Conv(
            in_channels,
            in_channels,
            kernel_size,
            stride,
            padding,
            dilation,
            in_channels,
            bias,
        )
        self.pointwise_conv = Conv(in_channels, out_channels, 1, 1, 0, 1, 1, bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，形状取决于卷积类型：
              - 1D: [N, C, T]
              - 2D: [N, C, H, W]
              - 3D: [N, C, D, H, W]

        Returns
        ----------
        - `torch.Tensor`
            - 输出张量，形状为 [N, out_channels, ...]
        """
        x = self.depthwise_conv(x)
        x = self.pointwise_conv(x)
        return x


class MaxNormConv(base.BaseModule):
    """最大范数约束卷积层

    在每次前向传播时对卷积层的权重进行最大范数约束，
    防止权重过大导致的梯度爆炸问题，提高训练稳定性。

    特别适用于 EEG 信号分析中的空间滤波器学习，
    可以防止某些电极权重过大而导致的过拟合。

    Examples
    ----------
    ```py
    # EEG 空间滤波器
    spatial_conv = nn.Conv1d(64, 32, kernel_size=1)
    max_norm_conv = MaxNormConv(conv=spatial_conv, max_norm=1.0)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    filtered_eeg = max_norm_conv(eeg_data)  # [N, 32, 1000]

    # 2D 卷积的最大范数约束
    conv2d = nn.Conv2d(128, 256, kernel_size=3, padding=1)
    max_norm_conv2d = MaxNormConv(conv=conv2d, max_norm=2.0)
    features = torch.randn(16, 128, 32, 32)
    output = max_norm_conv2d(features)
    ```
    """

    def __init__(self, conv, max_norm: float = 1.0):
        """初始化 MaxNormConv 层

        Parameters
        ----------
        - `conv` : `base.Conv`
            - 要应用最大范数约束的卷积层
        - `max_norm` : `float`, optional
            - 最大范数值, by default 1.0
        """
        super(MaxNormConv, self).__init__()
        self.conv = conv
        self.max_norm = max_norm

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播

        在每次前向传播前对卷积权重进行最大范数约束。

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，形状取决于卷积类型：
              - 1D: [N, C, T]
              - 2D: [N, C, H, W]
              - 3D: [N, C, D, H, W]

        Returns
        ----------
        - `torch.Tensor`
            - 卷积输出，权重已应用最大范数约束
        """
        # 应用最大范数约束到权重
        self.conv.weight.data = torch.renorm(
            self.conv.weight.data, p=2, dim=0, maxnorm=self.max_norm
        )
        return self.conv.forward(x)


class SincConv1d(base.BaseModule):
    """可学习的带通滤波器卷积层，专用于信号处理

    基于 Sinc 函数的卷积层，能够学习特定频段的滤波器参数，
    特别适用于语音和脑电信号处理。相比传统卷积，SincConv 能够
    学习更有意义的频域特征，提高信号分类和分析的性能。

    References
    ----------
    - **Paper**: Ravanelli, M., & Bengio, Y. (2018). "Speaker Recognition from Raw Waveform with SincNet."
      *Proceedings of SLT 2018*. https://arxiv.org/abs/1808.00158
    - **Official Implementation**: https://github.com/mravanelli/SincNet
    - **PyTorch Implementation**: https://github.com/pyannote/pyannote-audio
    - **EEG Application**: Borra, D., et al. (2021). "A Lightweight Multi-Scale Convolutional Neural Network for P300 Detection."

    Examples
    ----------
    ```py
    # EEG 信号滤波 (250Hz 采样率)
    sinc_conv = SincConv1d(in_channels=64, out_channels=32, kernel_size=251, sample_rate=250)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    filtered = sinc_conv(eeg_data)

    # 检查学习到的频率范围
    print(f"Low cutoff frequencies: {sinc_conv.low_hz}")
    print(f"High cutoff frequencies: {sinc_conv.high_hz}")

    # 语音信号处理 (16kHz 采样率)
    sinc_conv_audio = SincConv1d(in_channels=1, out_channels=80, kernel_size=251, sample_rate=16000)
    audio_data = torch.randn(16, 1, 16000)  # 1秒音频
    features = sinc_conv_audio(audio_data)
    ```
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        sample_rate: int = 250,
        min_low_hz: float = 50,
        min_band_hz: float = 50,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        bias: bool = False,
    ):
        """初始化 SincConv1d 层

        Parameters
        ----------
        - `in_channels` : `int`
            - 输入通道数
        - `out_channels` : `int`
            - 输出通道数（滤波器数量）
        - `kernel_size` : `int`
            - 卷积核大小，必须为奇数
        - `sample_rate` : `int`, optional
            - 采样率，用于频率计算, by default 250
        - `min_low_hz` : `float`, optional
            - 最小低截止频率, by default 50
        - `min_band_hz` : `float`, optional
            - 最小带宽, by default 50
        - `stride` : `int`, optional
            - 步长, by default 1
        - `padding` : `int`, optional
            - 填充, by default 0
        - `dilation` : `int`, optional
            - 膨胀, by default 1
        - `bias` : `bool`, optional
            - 是否使用偏置, by default False
        """
        super().__init__()

        if kernel_size % 2 == 0:
            raise ValueError('kernel_size must be odd')

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.sample_rate = sample_rate
        self.min_low_hz = min_low_hz
        self.min_band_hz = min_band_hz
        self.stride = stride
        self.padding = padding
        self.dilation = dilation

        # 频率参数（可学习）
        self.low_hz_ = nn.Parameter(
            torch.Tensor(out_channels).uniform_(min_low_hz, sample_rate / 2 - min_band_hz)
        )
        self.band_hz_ = nn.Parameter(
            torch.Tensor(out_channels).uniform_(min_band_hz, sample_rate / 2 - min_low_hz)
        )

        # 窗函数
        n_lin = torch.linspace(0, (kernel_size / 2) - 1, steps=int((kernel_size / 2)))
        self.window_ = 0.54 - 0.46 * torch.cos(2 * torch.pi * n_lin / kernel_size)

        # 时间轴
        n = (kernel_size - 1) / 2.0
        self.n_ = 2 * torch.pi * torch.arange(-n, 0).view(1, -1) / sample_rate

        if bias:
            self.bias = nn.Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)

        self.init_weights()

    def init_weights(self):
        """初始化权重"""
        if self.bias is not None:
            nn.init.constant_(self.bias, 0.0)

    @property
    def low_hz(self) -> torch.Tensor:
        """获取低截止频率"""
        return torch.abs(self.low_hz_) + self.min_low_hz

    @property
    def high_hz(self) -> torch.Tensor:
        """获取高截止频率"""
        return torch.clamp(
            self.low_hz + torch.abs(self.band_hz_) + self.min_band_hz,
            self.min_low_hz,
            self.sample_rate / 2,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入信号，形状为 [N, C, T]

        Returns
        ----------
        - `torch.Tensor`
            - 滤波后的信号
        """
        # 移动设备
        n_ = self.n_.to(x.device)
        window_ = self.window_.to(x.device)

        low = self.low_hz.view(-1, 1)
        high = self.high_hz.view(-1, 1)

        # 计算 sinc 滤波器
        f_times_t_low = torch.matmul(low, n_)
        f_times_t_high = torch.matmul(high, n_)

        band_pass_left = (
            (torch.sin(f_times_t_high) - torch.sin(f_times_t_low)) / (n_ / 2)
        ) * window_
        band_pass_center = 2 * (high - low)
        band_pass_right = torch.flip(band_pass_left, dims=[1])

        band_pass = torch.cat([band_pass_left, band_pass_center, band_pass_right], dim=1)
        band_pass = band_pass / (2 * high)

        # 重塑滤波器以适应标准卷积
        # 每个输出通道对应一个独立的带通滤波器
        filters = band_pass.view(self.out_channels, 1, self.kernel_size)

        # 扩展到所有输入通道 - 每个输出滤波器应用到所有输入通道
        filters = filters.repeat(1, self.in_channels, 1)

        return nn.functional.conv1d(
            x,
            filters,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
            dilation=self.dilation,
            groups=1,
        )

    def extra_repr(self) -> str:
        return (
            f'in_channels={self.in_channels}, out_channels={self.out_channels}, '
            f'kernel_size={self.kernel_size}, sample_rate={self.sample_rate}'
        )


class MultiScaleConv1d(base.BaseModule):
    """多尺度时序卷积模块

    通过并行的不同扩张率卷积捕获多个时间尺度的特征，
    特别适用于时序信号分析，如脑电信号的多频段特征提取。

    References
    ----------
    - **Paper**: Oord, A. v. d., et al. (2016). "WaveNet: A Generative Model for Raw Audio."
      *ICML 2016*. https://arxiv.org/abs/1609.03499
    - **Temporal Convolutional Networks**: Bai, S., et al. (2018). "An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling."
      https://arxiv.org/abs/1803.01271
    - **EEG Application**: Lawhern, V. J., et al. (2018). "EEGNet: a compact convolutional neural network for EEG-based brain–computer interfaces."

    Examples
    ----------
    ```py
    # EEG 多尺度特征提取
    multi_conv = MultiScaleConv1d(in_channels=64, out_channels=128, kernel_size=3, dilations=[1, 2, 4, 8])
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    multi_features = multi_conv(eeg_data)

    # 自定义扩张率
    custom_conv = MultiScaleConv1d(in_channels=32, out_channels=64, kernel_size=5, dilations=[1, 3, 9, 27])
    features = custom_conv(torch.randn(16, 32, 500))
    ```
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dilations: typing.Sequence = [1, 2, 4, 8],
        activation: str = 'relu',
        normalization: bool = True,
    ):
        """初始化 MultiScaleConv1d 层

        Parameters
        ----------
        - `in_channels` : `int`
            - 输入通道数
        - `out_channels` : `int`
            - 输出通道数
        - `kernel_size` : `int`
            - 卷积核大小
        - `dilations` : `list`, optional
            - 扩张率列表, by default [1, 2, 4, 8]
        - `activation` : `str`, optional
            - 激活函数类型, by default 'relu'
        - `normalization` : `bool`, optional
            - 是否使用批归一化, by default True
        """
        super().__init__()

        self.dilations = dilations
        self.num_scales = len(dilations)

        # 每个尺度的输出通道数
        channels_per_scale = out_channels // self.num_scales
        remaining_channels = out_channels % self.num_scales

        self.convs = nn.ModuleList()
        self.norms = nn.ModuleList() if normalization else None

        for i, dilation in enumerate(dilations):
            # 分配通道数
            scale_out_channels = channels_per_scale + (1 if i < remaining_channels else 0)

            # 计算padding以保持时序长度
            padding = (kernel_size - 1) * dilation // 2

            conv = nn.Conv1d(
                in_channels,
                scale_out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
                bias=not normalization,
            )
            self.convs.append(conv)

            if normalization:
                assert self.norms is not None
                self.norms.append(nn.BatchNorm1d(scale_out_channels))

        # 激活函数
        if activation == 'relu':
            self.activation = nn.ReLU(inplace=True)
        elif activation == 'gelu':
            self.activation = nn.GELU()
        elif activation == 'swish':
            self.activation = nn.SiLU()
        else:
            self.activation = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入信号，形状为 [N, C, T]

        Returns
        ----------
        - `torch.Tensor`
            - 多尺度特征，形状为 [N, out_channels, T]
        """
        outputs = []

        for i, conv in enumerate(self.convs):
            out = conv(x)

            if self.norms is not None:
                out = self.norms[i](out)

            out = self.activation(out)
            outputs.append(out)

        # 拼接所有尺度的特征
        return torch.cat(outputs, dim=1)

    def extra_repr(self) -> str:
        return f'dilations={self.dilations}, num_scales={self.num_scales}'
