"""EEG 专用注意力机制模块

实现专门针对脑电信号分析的注意力机制，包括通道注意力、时序注意力等。
"""

import typing

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from ..base import BaseModule
from .CBAM import ChannelAttention1D as BaseChannelAttention
from .CBAM import SpatialAttention1D as BaseSpatialAttention


class EEGChannelAttention(BaseModule):
    """EEG 通道注意力机制

    学习不同 EEG 电极通道的重要性权重

    Examples
    ----------
    ```py
    # 基本使用 - 64 通道 EEG
    channel_attn = EEGChannelAttention(num_channels=64)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    attended_eeg = channel_attn(eeg_data)

    # 查看通道权重
    weights = channel_attn.get_channel_weights(eeg_data)
    print(f"Channel attention weights shape: {weights.shape}")  # [N, C, 1]

    # 使用不同的池化策略
    spatial_attn = EEGChannelAttention(num_channels=128, pooling='both', reduction=8)
    eeg_128 = torch.randn(16, 128, 500)
    output = spatial_attn(eeg_128)
    ```
    """

    def __init__(
        self,
        num_channels: int,
        reduction: int = 16,
        pooling: str = 'avg',
        activation: str = 'relu',
    ):
        """初始化 EEGChannelAttention 层

        Parameters
        ----------
        - `num_channels` : `int`
            - EEG 通道数（电极数量）
        - `reduction` : `int`, optional
            - 降维比例，用于减少参数量, by default 16
        - `pooling` : `str`, optional
            - 池化策略: 'avg', 'max', 'both', by default 'avg'
        - `activation` : `str`, optional
            - 激活函数类型: 'relu', 'gelu', 'swish', by default 'relu'
        """
        super().__init__()

        self.num_channels = num_channels
        self.reduction = reduction
        self.pooling = pooling
        self.activation = activation

        # 直接使用增强后的 CBAM.ChannelAttention
        self.channel_attention = BaseChannelAttention(
            in_ch=num_channels,
            ratio=reduction,
            pooling=pooling,
            activation=activation,
            use_linear=True,  # EEG 使用 Linear 实现
        )

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

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

        Returns
        ----------
        - `torch.Tensor`
            - 通道注意力加权后的信号，形状为 [N, C, T]
        """
        attention_weights = self.channel_attention(x)  # [N, C, 1]
        return x * attention_weights

    def get_channel_weights(self, x: torch.Tensor) -> torch.Tensor:
        """获取通道注意力权重

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

        Returns
        ----------
        - `torch.Tensor`
            - 通道注意力权重，形状为 [N, C, 1]
        """
        return self.channel_attention(x)

    def extra_repr(self) -> str:
        return f'num_channels={self.num_channels}, reduction={self.reduction}, pooling={self.pooling}, activation={self.activation}'


class TemporalAttention(BaseModule):
    """时序注意力机制

    学习时间序列中不同时间点的重要性，特别适用于事件相关电位（ERP）
    分析和时间相关的脑电特征提取。

    Examples
    ----------
    ```py
    # 基本时序注意力
    temporal_attn = TemporalAttention(input_dim=64)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    attended_eeg = temporal_attn(eeg_data)

    # 获取时间注意力权重
    weights = temporal_attn.get_temporal_weights(eeg_data)
    print(f"Temporal attention weights shape: {weights.shape}")  # [N, 1, T]

    # 使用多头注意力
    multi_head_attn = TemporalAttention(input_dim=128, num_heads=8, use_multihead=True)
    output = multi_head_attn(torch.randn(16, 128, 500))
    ```
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dim: Optional[int] = None,
        num_heads: int = 1,
        use_multihead: bool = False,
        dropout: float = 0.1,
        kernel_size: int = 7,
    ):
        """初始化 TemporalAttention 层

        Parameters
        ----------
        - `input_dim` : `int`
            - 输入特征维度（通道数）
        - `hidden_dim` : `int`, optional
            - 隐藏层维度, by default None (使用 input_dim)
        - `num_heads` : `int`, optional
            - 多头注意力的头数, by default 1
        - `use_multihead` : `bool`, optional
            - 是否使用多头注意力, by default False
        - `dropout` : `float`, optional
            - Dropout 概率, by default 0.1
        - `kernel_size` : `int`, optional
            - 简单模式下的卷积核大小, by default 7
        """
        super().__init__()

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim or input_dim
        self.num_heads = num_heads
        self.use_multihead = use_multihead
        self.kernel_size = kernel_size

        if use_multihead:
            # 使用多头注意力实现
            self.multihead_attn = nn.MultiheadAttention(
                embed_dim=input_dim,
                num_heads=num_heads,
                dropout=dropout,
                batch_first=True,
            )
            self.dropout = nn.Dropout(dropout)
            self.layer_norm = nn.LayerNorm(input_dim)
        else:
            # 使用组合模式：复用 CBAM.SpatialAttention
            self.spatial_attention = BaseSpatialAttention(kernel_size=kernel_size)

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

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

        Returns
        ----------
        - `torch.Tensor`
            - 时序注意力加权后的信号，形状为 [N, C, T]
        """
        if self.use_multihead:
            # 使用多头注意力实现
            # 转换为 [N, T, C] 用于注意力计算
            x_transposed = x.transpose(1, 2)  # [N, T, C]

            attended, _ = self.multihead_attn(x_transposed, x_transposed, x_transposed)
            attended = self.dropout(attended)
            attended = self.layer_norm(attended + x_transposed)

            # 转换回 [N, C, T]
            return attended.transpose(1, 2)
        else:
            # 使用组合模式：复用 CBAM.SpatialAttention
            # 对于 1D 信号，空间注意力就是时序注意力
            attention_weights = self.spatial_attention(x)  # [N, 1, T]
            return x * attention_weights

    def get_temporal_weights(self, x: torch.Tensor) -> torch.Tensor:
        """获取时序注意力权重

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

        Returns
        ----------
        - `torch.Tensor`
            - 时序泣意力权重，形状为 [N, 1, T] 或 [N, T, T]
        """
        if self.use_multihead:
            raise NotImplementedError(
                'Weight extraction not supported for multihead attention'
            )
        else:
            # 使用组合模式：复用 CBAM.SpatialAttention
            return self.spatial_attention(x)  # [N, 1, T]

    def extra_repr(self) -> str:
        base_info = f'input_dim={self.input_dim}, use_multihead={self.use_multihead}'
        if self.use_multihead:
            base_info += f', hidden_dim={self.hidden_dim}, num_heads={self.num_heads}'
        else:
            base_info += f', kernel_size={self.kernel_size}'
        return base_info


class SpectralAttention(BaseModule):
    """频谱注意力机制

    在频域中学习不同频率成分的重要性，特别适用于脑电信号的
    频域分析和频段特征提取。

    References
    ----------
    - **Paper**: Hu, J., et al. (2018). "Squeeze-and-Excitation Networks."
      *CVPR 2018*. https://arxiv.org/abs/1709.01507
    - **Frequency Domain**: Welch, P. (1967). "The use of fast Fourier transform for the estimation of power spectra."
    - **EEG Spectral Analysis**: Sanei, S., & Chambers, J. A. (2013). "EEG signal processing."

    Examples
    ----------
    ```py
    # 基本频谱注意力
    spectral_attn = SpectralAttention(num_channels=64, fft_size=256)
    eeg_data = torch.randn(32, 64, 1000)  # [N, C, T]
    attended_eeg = spectral_attn(eeg_data)

    # 自定义频率范围
    custom_attn = SpectralAttention(
        num_channels=128,
        fft_size=512,
        freq_bands=[(0.5, 4), (4, 8), (8, 13), (13, 30)]  # δ, θ, α, β 波段
    )
    output = custom_attn(torch.randn(16, 128, 2000))
    ```
    """

    def __init__(
        self,
        num_channels: int,
        fft_size: int = 256,
        hop_length: Optional[int] = None,
        freq_bands: Optional[typing.Sequence] = None,
        sample_rate: int = 250,
        reduction: int = 16,
    ):
        """初始化 SpectralAttention 层

        Parameters
        ----------
        - `num_channels` : `int`
            - 输入通道数
        - `fft_size` : `int`, optional
            - FFT 窗口大小, by default 256
        - `hop_length` : `int`, optional
            - 跳跃长度, by default None (fft_size // 4)
        - `freq_bands` : `list`, optional
            - 关注的频率段列表 [(low, high), ...], by default None
        - `sample_rate` : `int`, optional
            - 采样率, by default 250
        - `reduction` : `int`, optional
            - 降维比例, by default 16
        """
        super().__init__()

        self.num_channels = num_channels
        self.fft_size = fft_size
        self.hop_length = hop_length or fft_size // 4
        self.freq_bands = freq_bands
        self.sample_rate = sample_rate

        # 频谱特征维度
        self.freq_bins = fft_size // 2 + 1

        if freq_bands:
            # 使用指定频段
            self.band_indices = self._get_band_indices(freq_bands, sample_rate, fft_size)
            feature_dim = len(freq_bands) * num_channels
        else:
            # 使用全频谱
            feature_dim = self.freq_bins * num_channels

        # 注意力网络
        hidden_dim = max(1, feature_dim // reduction)
        self.attention_net = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, num_channels),
            nn.Sigmoid(),
        )

    def _get_band_indices(
        self, freq_bands: typing.Sequence, sample_rate: int, fft_size: int
    ) -> list:
        """获取频段对应的FFT索引"""
        indices = []
        for low_freq, high_freq in freq_bands:
            low_idx = int(low_freq * fft_size / sample_rate)
            high_idx = int(high_freq * fft_size / sample_rate)
            indices.append((low_idx, high_idx))
        return indices

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

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

        Returns
        ----------
        - `torch.Tensor`
            - 频谱注意力加权后的信号，形状为 [N, C, T]
        """
        batch_size, num_channels, seq_len = x.size()

        # 计算短时傅里叶变换
        spectrograms = []
        for i in range(num_channels):
            stft = torch.stft(
                x[:, i, :],
                n_fft=self.fft_size,
                hop_length=self.hop_length,
                return_complex=True,
            )
            power_spec = torch.abs(stft) ** 2  # [N, freq_bins, time_frames]
            spectrograms.append(power_spec)

        # 堆叠所有通道的频谱
        spectrograms = torch.stack(spectrograms, dim=1)  # [N, C, freq_bins, time_frames]

        # 提取频谱特征
        if self.freq_bands:
            # 使用指定频段
            band_features = []
            for low_idx, high_idx in self.band_indices:
                band_power = torch.mean(
                    spectrograms[:, :, low_idx:high_idx, :], dim=(2, 3)
                )  # [N, C]
                band_features.append(band_power)
            spectral_features = torch.cat(band_features, dim=1)  # [N, C * num_bands]
        else:
            # 使用全频谱
            spectral_features = torch.mean(spectrograms, dim=3)  # [N, C, freq_bins]
            spectral_features = spectral_features.view(
                batch_size, -1
            )  # [N, C * freq_bins]

        # 计算通道注意力权重
        attention_weights = self.attention_net(spectral_features)  # [N, C]
        attention_weights = attention_weights.unsqueeze(-1)  # [N, C, 1]

        # 应用注意力权重
        return x * attention_weights

    def extra_repr(self) -> str:
        return (
            f'num_channels={self.num_channels}, fft_size={self.fft_size}, '
            f'sample_rate={self.sample_rate}'
        )
