"""Squeeze-and-Excitation 模块

实现 1D 版本的 SE 注意力机制，用于自适应地重新校准通道特征响应。
"""

import torch
from torch import nn

from .base import BaseModule


class SEBlock(BaseModule):
    """Squeeze-and-Excitation Block for 1D signals

    通过全局平均池化和两个全连接层学习通道间的依赖关系，
    为每个通道生成注意力权重，提升重要通道的特征表达。

    References
    ----------
    - **Paper**: Hu, J., Shen, L., & Sun, G. (2018). "Squeeze-and-Excitation Networks."
      *CVPR 2018*. https://arxiv.org/abs/1709.01507
    - **Official Implementation**: https://github.com/hujie-frank/SENet

    Examples
    ----------
    ```py
    # EEG 信号通道注意力
    se_block = SEBlock(in_channels=64, reduction=16)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    enhanced_eeg = se_block(eeg_data)  # [N, C, T]

    # 音频信号处理
    audio_se = SEBlock(in_channels=128, reduction=8)
    audio_features = torch.randn(16, 128, 2048)
    output = audio_se(audio_features)
    ```
    """

    def __init__(self, in_channels: int, reduction: int = 16):
        """初始化 SEBlock 层

        Parameters
        ----------
        - `in_channels` : `int`
            - 输入通道数
        - `reduction` : `int`, optional
            - 降维比例，控制中间层的通道数, by default 16
        """
        super(SEBlock, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction, in_channels, bias=False),
            nn.Sigmoid(),
        )

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

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

        Returns
        ----------
        - `torch.Tensor`
            - SE 注意力加权后的信号，形状为 [N, C, T]
        """
        b, c, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1)
        return x * y.expand_as(x)


class SEBlock2D(BaseModule):
    """Squeeze-and-Excitation Block for 2D signals

    2D 版本的 SE 注意力机制，适用于图像或 2D 时频表示的信号处理。

    References
    ----------
    - **Paper**: Hu, J., Shen, L., & Sun, G. (2018). "Squeeze-and-Excitation Networks."
      *CVPR 2018*. https://arxiv.org/abs/1709.01507

    Examples
    ----------
    ```py
    # 图像特征增强
    se_block_2d = SEBlock2D(in_channels=256, reduction=16)
    image_features = torch.randn(32, 256, 28, 28)  # [N, C, H, W]
    enhanced_features = se_block_2d(image_features)  # [N, C, H, W]

    # EEG 时频图处理
    eeg_se = SEBlock2D(in_channels=64, reduction=8)
    spectrogram = torch.randn(16, 64, 129, 100)  # [N, C, F, T]
    output = eeg_se(spectrogram)
    ```
    """

    def __init__(self, in_channels: int, reduction: int = 16):
        """初始化 SEBlock2D 层

        Parameters
        ----------
        - `in_channels` : `int`
            - 输入通道数
        - `reduction` : `int`, optional
            - 降维比例，控制中间层的通道数, by default 16
        """
        super(SEBlock2D, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction, in_channels, bias=False),
            nn.Sigmoid(),
        )

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

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

        Returns
        ----------
        - `torch.Tensor`
            - SE 注意力加权后的信号，形状为 [N, C, H, W]
        """
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)
