"""CBAM (Convolutional Block Attention Module) 注意力机制

实现通道注意力和空间注意力的组合，用于增强重要特征并抑制无关信息。
特别适用于 1D 信号处理，如 EEG、音频等时序数据。
"""

import torch
import torch.nn as nn

from zyplib.nn.modules.base import BaseModule


class ChannelAttention1D(BaseModule):
    """通道注意力模块

    通过全局平均池化和最大池化捕获通道间的依赖关系，
    学习每个通道的重要性权重。

    References
    ----------
    - **Paper**: Woo, S., et al. (2018). "CBAM: Convolutional Block Attention Module."
      *ECCV 2018*. https://arxiv.org/abs/1807.06521

    Examples
    ----------
    ```py
    # EEG 通道注意力
    channel_attn = ChannelAttention(in_ch=64, ratio=16)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    attention_weights = channel_attn(eeg_data)  # [N, C, 1]
    attended_eeg = eeg_data * attention_weights
    ```
    """

    def __init__(
        self,
        in_ch: int,
        ratio: int = 16,
        pooling: str = 'both',
        activation: str = 'relu',
        use_linear: bool = False,
    ):
        """初始化 ChannelAttention 层

        Parameters
        ----------
        - `in_ch` : `int`
            - 输入通道数
        - `ratio` : `int`, optional
            - 降维比例，控制中间层通道数, by default 16
        - `pooling` : `str`, optional
            - 池化策略: 'avg', 'max', 'both', by default 'both'
        - `activation` : `str`, optional
            - 激活函数类型: 'relu', 'gelu', 'swish', by default 'relu'
        - `use_linear` : `bool`, optional
            - 是否使用 Linear 层而非 Conv1d, by default False
        """
        super(ChannelAttention1D, self).__init__()

        self.in_ch = in_ch
        self.ratio = ratio
        self.pooling = pooling
        self.activation = activation
        self.use_linear = use_linear

        # 池化层
        self.avg_pool = nn.AdaptiveAvgPool1d(1)
        self.max_pool = nn.AdaptiveMaxPool1d(1)

        # 计算输入维度
        if pooling == 'both':
            input_dim = in_ch * 2
        else:
            input_dim = in_ch

        hidden_dim = max(1, in_ch // ratio)

        # 根据配置选择实现方式
        if use_linear:
            self.fc = nn.Sequential(
                nn.Linear(input_dim, hidden_dim, bias=False),
                self._get_activation(activation),
                nn.Linear(hidden_dim, in_ch, bias=False),
            )
        else:
            # 传统 Conv1d 实现（仅支持 pooling='both'）
            if pooling != 'both':
                raise ValueError(
                    f"Conv1d mode only supports pooling='both', got '{pooling}'"
                )
            self.fc = nn.Sequential(
                nn.Conv1d(in_ch, hidden_dim, 1, bias=False),
                self._get_activation(activation),
                nn.Conv1d(hidden_dim, in_ch, 1, bias=False),
            )

        self.sigmoid = nn.Sigmoid()

    def _get_activation(self, activation: str) -> nn.Module:
        """获取激活函数"""
        if activation == 'relu':
            return nn.ReLU(inplace=True)
        elif activation == 'gelu':
            return nn.GELU()
        elif activation == 'swish':
            return nn.SiLU()
        else:
            return nn.ReLU(inplace=True)

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

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

        Returns
        ----------
        - `torch.Tensor`
            - 通道注意力权重，形状为 [N, C, 1]
        """
        if self.use_linear:
            # Linear 实现：支持灵活的池化策略
            batch_size, num_channels, seq_len = x.size()

            # 收集池化特征
            features = []
            if self.pooling in ['avg', 'both']:
                avg_out = self.avg_pool(x).view(batch_size, num_channels)  # [N, C]
                features.append(avg_out)
            if self.pooling in ['max', 'both']:
                max_out = self.max_pool(x).view(batch_size, num_channels)  # [N, C]
                features.append(max_out)

            # 拼接特征
            if len(features) > 1:
                pooled = torch.cat(features, dim=1)  # [N, C*2] if both
            else:
                pooled = features[0]  # [N, C]

            # 计算注意力权重
            attention_weights = self.fc(pooled)  # [N, C]
            return self.sigmoid(attention_weights).unsqueeze(-1)  # [N, C, 1]
        else:
            # 传统 Conv1d 实现（仅支持 pooling='both'）
            avg_out = self.fc(self.avg_pool(x))
            max_out = self.fc(self.max_pool(x))
            out = avg_out + max_out
            return self.sigmoid(out)


class SpatialAttention1D(BaseModule):
    """空间注意力模块

    通过沿通道维度的平均池化和最大池化，学习空间位置的重要性。
    对于 1D 信号，这相当于学习时间维度上不同位置的注意力权重。

    References
    ----------
    - **Paper**: Woo, S., et al. (2018). "CBAM: Convolutional Block Attention Module."
      *ECCV 2018*. https://arxiv.org/abs/1807.06521

    Examples
    ----------
    ```py
    # EEG 时间注意力
    spatial_attn = SpatialAttention1D(kernel_size=7)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    attention_weights = spatial_attn(eeg_data)  # [N, 1, T]
    attended_eeg = eeg_data * attention_weights
    ```
    """

    def __init__(self, kernel_size: int = 7):
        """初始化 SpatialAttention 层

        Parameters
        ----------
        - `kernel_size` : `int`, optional
            - 卷积核大小，用于空间特征提取, by default 7
        """
        super(SpatialAttention1D, self).__init__()
        assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'

        self.conv1 = nn.Conv1d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

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

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

        Returns
        ----------
        - `torch.Tensor`
            - 空间注意力权重，形状为 [N, 1, T]
        """
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)


class CBAM1D(BaseModule):
    """Convolutional Block Attention Module

    结合通道注意力和空间注意力的完整 CBAM 模块。
    先应用通道注意力，再应用空间注意力，实现双重注意力机制。

    References
    ----------
    - **Paper**: Woo, S., et al. (2018). "CBAM: Convolutional Block Attention Module."
      *ECCV 2018*. https://arxiv.org/abs/1807.06521
    - **Official Implementation**: https://github.com/Jongchan/attention-module

    Examples
    ----------
    ```py
    # EEG 信号双重注意力
    cbam = CBAM(in_ch=64, ratio=16, kernel_size=7)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    enhanced_eeg = cbam(eeg_data)  # [N, C, T]

    # 音频特征增强
    audio_cbam = CBAM(in_ch=128, ratio=8, kernel_size=5)
    audio_features = torch.randn(16, 128, 2048)
    output = audio_cbam(audio_features)
    ```
    """

    def __init__(self, in_ch: int, ratio: int = 16, kernel_size: int = 7):
        """初始化 CBAM 层

        Parameters
        ----------
        - `in_ch` : `int`
            - 输入通道数
        - `ratio` : `int`, optional
            - 通道注意力的降维比例, by default 16
        - `kernel_size` : `int`, optional
            - 空间注意力的卷积核大小, by default 7
        """
        super(CBAM1D, self).__init__()
        self.ca = ChannelAttention1D(in_ch, ratio)
        self.sa = SpatialAttention1D(kernel_size)

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

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

        Returns
        ----------
        - `torch.Tensor`
            - CBAM 注意力增强后的信号，形状为 [N, C, T]
        """
        x = x * self.ca(x)
        x = x * self.sa(x)
        return x
