import math
from typing import Optional, cast

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

from zyplib.nn.modules.base import BaseModule

from .attn_utils import combine_heads, split_heads


class SlidingWindowAttention(BaseModule):
    """滑动窗口注意力机制 | 一种 Local Attention

    通过限制每个 token 只关注其局部窗口内的 token，将注意力的时间复杂度从 O(n^2) 降低到 O(n*w)，
    其中 n 是序列长度，w 是窗口大小。这对于处理具有强局部相关性的长序列（如脑电信号）非常有效。

    References
    ----------
    - **Paper**: Beltagy, I., et al. (2020). "Longformer: The Long-Document Transformer."
      *ArXiv preprint arXiv:2004.05150*.

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.efficient import SlidingWindowAttention
    >>>
    >>> # 处理脑电信号序列
    >>> seq_len, embed_dim, window_size = 2048, 128, 64
    >>> model = SlidingWindowAttention(embed_dim=embed_dim, num_heads=8, window_size=window_size)
    >>> x = torch.randn(16, seq_len, embed_dim)
    >>>
    >>> output, attn_weights = model(x, x, x)
    >>> print(f"输入形状: {x.shape}")
    >>> print(f"输出形状: {output.shape}")
    >>> print(f"注意力权重形状: {attn_weights.shape}")
    输入形状: torch.Size([16, 2048, 128])
    输出形状: torch.Size([16, 2048, 128])
    注意力权重形状: torch.Size([16, 8, 2048, 2048])
    ```
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        window_size: int,
        dropout: float = 0.1,
        memory_efficient: bool = False,
    ):
        """初始化 SlidingWindowAttention 层

        Parameters
        ----------
        - `embed_dim` : `int`
            - 输入特征的维度
        - `num_heads` : `int`
            - 注意力头的数量
        - `window_size` : `int`
            - 滑动窗口的大小，建议使用奇数以保持对称性
        - `dropout` : `float`, optional
            - Dropout 比例, by default 0.1
        - `memory_efficient` : `bool`, optional
            - 是否使用内存高效的实现，避免创建完整的 O(n²) 矩阵, by default True
        """
        super().__init__()
        assert embed_dim % num_heads == 0, 'embed_dim 必须能被 num_heads 整除'

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.window_size = window_size
        self.half_window = window_size // 2
        self.memory_efficient = memory_efficient

        # 线性投影层
        self.W_q = nn.Linear(embed_dim, embed_dim)
        self.W_k = nn.Linear(embed_dim, embed_dim)
        self.W_v = nn.Linear(embed_dim, embed_dim)
        self.fc = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)
        self.scale: float = self.head_dim**-0.5

    @staticmethod
    def create_sliding_window_mask(
        seq_len: int, window_size: int, device: torch.device
    ) -> torch.Tensor:
        """创建滑动窗口掩码

        Parameters
        ----------
        - `seq_len` : `int`
            - 序列长度
        - `window_size` : `int`
            - 窗口大小
        - `device` : `torch.device`
            - 设备类型

        Returns
        ----------
        - `torch.Tensor`
            - 形状为 (seq_len, seq_len) 的布尔掩码，True 表示允许关注

        Examples
        --------
        ```python
        >>> import torch
        >>> from zyplib.nn.modules.attention.attn_utils import create_sliding_window_mask
        >>>
        >>> mask = create_sliding_window_mask(10, 3, torch.device('cpu'))
        >>> print(mask.shape)  # torch.Size([10, 10])
        ```
        """
        mask = torch.zeros(seq_len, seq_len, dtype=torch.bool, device=device)
        half_window = window_size // 2

        # for i in range(seq_len):
        #     start = max(0, i - half_window)
        #     end = min(seq_len, i + half_window + 1)
        #     mask[i, start:end] = True

        # 使用广播操作，更高效
        i_indices = torch.arange(seq_len, device=device).unsqueeze(1)
        j_indices = torch.arange(seq_len, device=device).unsqueeze(0)

        # 计算距离并创建掩码
        distance = torch.abs(i_indices - j_indices)
        mask = distance <= half_window

        return mask

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """滑动窗口注意力的前向传播

        Parameters
        ----------
        - `q` : `torch.Tensor`
            - 查询张量，形状: (batch_size, seq_len_q, embed_dim)
        - `k` : `torch.Tensor`
            - 键张量，形状: (batch_size, seq_len_kv, embed_dim)
        - `v` : `torch.Tensor`
            - 值张量，形状: (batch_size, seq_len_kv, embed_dim)
        - `mask` : `Optional[torch.Tensor]`, optional
            - 额外的注意力掩码, by default None

        Returns
        ----------
        - `tuple[torch.Tensor, torch.Tensor]`
            - 输出张量和注意力权重
        """
        batch_size, seq_len_q, _ = q.shape
        seq_len_kv = k.shape[1]

        # 线性投影
        q = self.W_q(q)
        k = self.W_k(k)
        v = self.W_v(v)

        # 拆分为多头
        q = split_heads(
            q, self.num_heads, self.head_dim
        )  # (batch, num_heads, seq_len_q, head_dim)
        k = split_heads(
            k, self.num_heads, self.head_dim
        )  # (batch, num_heads, seq_len_kv, head_dim)
        v = split_heads(
            v, self.num_heads, self.head_dim
        )  # (batch, num_heads, seq_len_kv, head_dim)

        if self.memory_efficient and seq_len_q == seq_len_kv:
            # 内存高效的实现：避免创建完整的 O(n²) 矩阵
            context, attn_weights = self._memory_efficient_forward(q, k, v, mask)
        else:
            # 标准实现：适用于较短序列或交叉注意力
            context, attn_weights = self._standard_forward(q, k, v, mask)

        # 合并多头
        context = combine_heads(context, self.embed_dim)

        # 最终线性投影
        output = self.fc(context)

        return output, attn_weights

    def _standard_forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """标准的滑动窗口注意力实现

        适用于较短序列，会创建完整的注意力矩阵
        """
        seq_len_q = q.shape[2]
        seq_len_kv = k.shape[2]

        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale

        # 创建滑动窗口掩码（假设 q 和 k 序列长度相同）
        if seq_len_q == seq_len_kv:
            window_mask = SlidingWindowAttention.create_sliding_window_mask(
                seq_len_q, self.window_size, q.device
            )
            # 将窗口外的位置设置为负无穷
            scores = scores.masked_fill(
                ~window_mask.unsqueeze(0).unsqueeze(0), float('-inf')
            )

        # 应用额外的掩码（如果提供）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # 计算注意力权重
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 应用注意力权重
        context = torch.matmul(attn_weights, v)

        return context, attn_weights

    def _memory_efficient_forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """内存高效的滑动窗口注意力实现

        通过分块计算避免创建完整的 O(n²) 注意力矩阵，
        内存复杂度降低到 O(n*w)，其中 w 是窗口大小
        """
        raise NotImplementedError


class LinformerAttention(BaseModule):
    """Linformer 注意力机制 | 一种 Approximated Attention

    通过在序列维度上应用线性投影，将自注意力的时间复杂度从 O(n^2) 降低到 O(n*k)，
    其中 n 是序列长度，k 是投影维度。这对于处理长序列非常有效。

    References
    ----------
    - **Paper**: Wang, S., et al. (2020). "Linformer: Self-Attention with Linear Complexity."
      *ArXiv preprint arXiv:2006.04768*.

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.efficient import LinformerAttention
    >>>
    >>> # 处理一个长序列
    >>> seq_len, embed_dim, proj_k = 4096, 128, 256
    >>> model = LinformerAttention(embed_dim=embed_dim, num_heads=8, proj_k=proj_k)
    >>> x = torch.randn(32, seq_len, embed_dim)
    >>>
    >>> output, attn_weights = model(x, x, x)
    >>> print(f"输入形状: {x.shape}")
    >>> print(f"输出形状: {output.shape}")
    >>> print(f"注意力权重形状: {attn_weights.shape}")
    输入形状: torch.Size([32, 4096, 128])
    输出形状: torch.Size([32, 4096, 128])
    注意力权重形状: torch.Size([32, 8, 4096, 256])
    ```
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        proj_k: int,
        max_seq_len: int,
        dropout: float = 0.1,
        bias: bool = True,
        share_kv: bool = True,  # 是否共享 K/V 投影
    ):
        """初始化 LinformerAttention 层

        Parameters
        ----------
        - `embed_dim` : `int`
            - 输入特征的维度。
        - `num_heads` : `int`
            - 注意力头的数量。
        - `proj_k` : `int`
            - 投影后的序列长度，即论文中的 k。
        - `dropout` : `float`, optional
            - Dropout 比例, by default 0.1.
        - `bias` : `bool`, optional
            - 是否在 W_q, W_k, W_v 中使用偏置, by default True.
        """
        super().__init__()
        if embed_dim % num_heads != 0:
            raise ValueError(
                f'embed_dim ({embed_dim}) 必须能被 num_heads ({num_heads}) 整除'
            )

        from .multihead import ScaledDotProductAttention

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.d_k = embed_dim // num_heads
        self.proj_k = proj_k
        self.max_seq_len = max_seq_len

        # 标准 Q/K/V 投影
        self.W_q = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.W_k = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.W_v = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.fc = nn.Linear(embed_dim, embed_dim)

        # Linformer 的核心：固定的投影矩阵
        if share_kv:
            # 共享 K/V 投影矩阵（推荐，减少参数）
            self.E = nn.Linear(max_seq_len, proj_k, bias=False)
            self.F = self.E  # 共享权重
        else:
            # 分别的 K/V 投影矩阵
            self.E = nn.Linear(max_seq_len, proj_k, bias=False)
            self.F = nn.Linear(max_seq_len, proj_k, bias=False)

        self.attention = ScaledDotProductAttention(dropout=dropout)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Linformer 的前向传播过程。

        与标准多头注意力的区别在于，它首先对 Key 和 Value 张量进行线性投影，
        以减小序列长度维度，然后再计算注意力。

        Parameters
        ----------
        - `q` : `torch.Tensor`
            - 查询张量, 形状: `(batch_size, seq_len_q, embed_dim)`.
        - `k` : `torch.Tensor`
            - 键张量, 形状: `(batch_size, seq_len_kv, embed_dim)`.
        - `v` : `torch.Tensor`
            - 值张量, 形状: `(batch_size, seq_len_kv, embed_dim)`.
        - `mask` : `Optional[torch.Tensor]`, optional
            - 注意力掩码, by default None.

        Returns
        ----------
        - `tuple[torch.Tensor, torch.Tensor]`
            - 输出张量和注意力权重。
        """
        seq_len = k.shape[1]

        # 检查序列长度限制
        if seq_len > self.max_seq_len:
            raise ValueError(f'序列长度 {seq_len} 超过最大值 {self.max_seq_len}')

        # 1. 标准 Q/K/V 投影
        q = self.W_q(q)
        k = self.W_k(k)
        v = self.W_v(v)

        # 2. Linformer 投影：在序列维度上投影 K/V
        if seq_len < self.max_seq_len:
            # 如果序列长度小于最大长度，需要填充或截取投影矩阵
            E_weight = self.E.weight[:, :seq_len]  # 截取相应大小
            F_weight = self.F.weight[:, :seq_len]
            k = F.linear(k.transpose(1, 2), E_weight).transpose(1, 2)
            v = F.linear(v.transpose(1, 2), F_weight).transpose(1, 2)
        else:
            # 标准投影
            k = self.E(k.transpose(1, 2)).transpose(1, 2)
            v = self.F(v.transpose(1, 2)).transpose(1, 2)

        # 3. 拆分多头
        q = split_heads(q, self.num_heads, self.d_k)
        k = split_heads(k, self.num_heads, self.d_k)
        v = split_heads(v, self.num_heads, self.d_k)

        # 4. 注意力计算（掩码需要相应调整）
        if mask is not None:
            # 简单策略：如果掩码尺寸不匹配就忽略
            if mask.shape[-1] != self.proj_k:
                mask = None

        context = self.attention(q, k, v, mask=mask)

        # 5. 合并头和最终投影
        context = combine_heads(context, self.embed_dim)
        output = self.fc(context)

        return output


class PerformerAttention(BaseModule):
    """Performer 注意力机制

    通过使用随机特征映射来近似 softmax 核函数，将注意力的时间复杂度从 O(n^2) 降低到 O(n)。
    这种方法特别适合处理非常长的序列，如长时程脑电信号分析。

    References
    ----------
    - **Paper**: Choromanski, K., et al. (2020). "Rethinking Attention with Performers."
    *ArXiv preprint arXiv:2009.14794*.

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.efficient import PerformerAttention
    >>>
    >>> # 处理超长脑电信号序列
    >>> seq_len, embed_dim, num_features = 8192, 128, 256
    >>> model = PerformerAttention(embed_dim=embed_dim, num_heads=8, num_features=num_features)
    >>> x = torch.randn(8, seq_len, embed_dim)
    >>>
    >>> output = model(x, x, x)
    >>> print(f"输入形状: {x.shape}")
    >>> print(f"输出形状: {output.shape}")
    输入形状: torch.Size([8, 8192, 128])
    输出形状: torch.Size([8, 8192, 128])
    ```
    """

    def __init__(
        self, embed_dim: int, num_heads: int, num_features: int, dropout: float = 0.1
    ):
        """初始化 PerformerAttention 层

        Parameters
        ----------
        - `embed_dim` : `int`
            - 输入特征的维度
        - `num_heads` : `int`
            - 注意力头的数量
        - `num_features` : `int`
            - 随机特征映射的维度，通常设置为 head_dim 的 2-4 倍
        - `dropout` : `float`, optional
            - Dropout 比例, by default 0.1
        """
        super().__init__()
        assert embed_dim % num_heads == 0, 'embed_dim 必须能被 num_heads 整除'

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.num_features = num_features

        # 线性投影层
        self.W_q = nn.Linear(embed_dim, embed_dim)
        self.W_k = nn.Linear(embed_dim, embed_dim)
        self.W_v = nn.Linear(embed_dim, embed_dim)
        self.fc = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)

        # 随机特征映射的投影矩阵（固定不变）
        self.register_buffer('projection_matrix', self._create_projection_matrix())

    def _create_projection_matrix(self) -> torch.Tensor:
        """创建随机特征映射的投影矩阵

        Returns
        ----------
        - `torch.Tensor`
            - 形状为 (head_dim, num_features) 的随机投影矩阵
        """
        nb_rows = self.num_features // 2  # 因为要 concat cos 和 sin
        device = next(self.parameters()).device if hasattr(self, '_parameters') else None
        return self.create_orthogonal_random_matrix(self.head_dim, nb_rows, device)

    @staticmethod
    def create_orthogonal_random_matrix(
        input_dim: int, output_dim: int, device: torch.device | None = None
    ) -> torch.Tensor:
        """创建随机投影矩阵（用于 Performer 等）

        Parameters
        ----------
        - `input_dim` : `int`
            - 输入维度
        - `output_dim` : `int`
            - 输出维度
        - `device` : `torch.device`, optional
            - 设备类型, by default None

        Returns
        ----------
        - `torch.Tensor`
            - 随机投影矩阵，形状: (input_dim, output_dim)
        """
        if input_dim >= output_dim:
            # 标准情况：可以直接QR分解
            unstructured_block = torch.randn(input_dim, output_dim, device=device)
            q, _ = torch.linalg.qr(unstructured_block, mode='reduced')
            return q
        else:
            # input_dim < output_dim：需要创建更大的矩阵
            size = max(input_dim, output_dim)
            unstructured_block = torch.randn(size, size, device=device)
            q, _ = torch.linalg.qr(unstructured_block, mode='reduced')
            return q[:input_dim, :output_dim]

    def _kernel_feature_creator(self, x: torch.Tensor) -> torch.Tensor:
        """创建核函数特征映射

        使用随机特征映射来近似 exp(x^T y) 核函数

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，形状: (..., head_dim)

        Returns
        ----------
        - `torch.Tensor`
            - 特征映射后的张量，形状: (..., num_features)
        """
        return self.kernel_feature_mapping(x, cast(torch.Tensor, self.projection_matrix))

    @staticmethod
    def kernel_feature_mapping(
        x: torch.Tensor, projection_matrix: torch.Tensor
    ) -> torch.Tensor:
        """核函数特征映射

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，形状: (..., input_dim)
        - `projection_matrix` : `torch.Tensor`
            - 投影矩阵，形状: (input_dim, num_features)

        Returns
        ----------
        - `torch.Tensor`
            - 特征映射后的张量，形状: (..., num_features)

        Examples
        --------
        ```python
        >>> import torch
        >>> from zyplib.nn.modules.attention.attn_utils import kernel_feature_mapping
        >>>
        >>> x = torch.randn(32, 8, 100, 64)
        >>> proj_matrix = torch.randn(64, 128)
        >>> features = kernel_feature_mapping(x, proj_matrix)
        >>> print(features.shape)  # torch.Size([32, 8, 100, 128])
        ```
        """
        """FAVOR+ 特征映射"""
        # 计算投影
        projections = torch.matmul(x, projection_matrix)

        # 数值稳定性处理
        data_normalizer = (x**2).sum(dim=-1, keepdim=True) / 2.0

        # FAVOR+ 特征映射：[exp(ω^T x - ||x||²/2) * cos(ω^T x), exp(ω^T x - ||x||²/2) * sin(ω^T x)]
        ratio = 1.0 / math.sqrt(projection_matrix.shape[0])

        exp_features = torch.exp(-data_normalizer + projections * ratio)
        cos_features = torch.cos(projections) * exp_features
        sin_features = torch.sin(projections) * exp_features

        return torch.cat([cos_features, sin_features], dim=-1)

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        """Performer 注意力的前向传播

        注意 Preformer 不计算完整注意力权重矩阵，所以不返回 attn_weights

        Parameters
        ----------
        - `q` : `torch.Tensor`
            - 查询张量，形状: (batch_size, seq_len_q, embed_dim)
        - `k` : `torch.Tensor`
            - 键张量，形状: (batch_size, seq_len_kv, embed_dim)
        - `v` : `torch.Tensor`
            - 值张量，形状: (batch_size, seq_len_kv, embed_dim)

        Returns
        ----------
        - **output** (`torch.Tensor`): 最终的输出张量, 形状: `(batch_size, seq_len_q, embed_dim)`.
        """
        # 线性投影
        q = self.W_q(q)
        k = self.W_k(k)
        v = self.W_v(v)

        # 拆分为多头
        q = split_heads(
            q, self.num_heads, self.head_dim
        )  # (batch, num_heads, seq_len_q, head_dim)
        k = split_heads(
            k, self.num_heads, self.head_dim
        )  # (batch, num_heads, seq_len_kv, head_dim)
        v = split_heads(
            v, self.num_heads, self.head_dim
        )  # (batch, num_heads, seq_len_kv, head_dim)

        # 应用核函数特征映射
        q_features: torch.Tensor = self._kernel_feature_creator(
            q
        )  # (batch, num_heads, seq_len_q, num_features)
        k_features: torch.Tensor = self._kernel_feature_creator(
            k
        )  # (batch, num_heads, seq_len_kv, num_features)

        # Performer 的线性注意力计算
        # 计算 K'^T @ V: (batch, num_heads, num_features, head_dim)
        kv = torch.matmul(k_features.transpose(-2, -1), v)

        # 计算 Q' @ (K'^T @ V): (batch, num_heads, seq_len_q, head_dim)
        context = torch.matmul(q_features, kv)

        # 归一化因子: Q' @ (sum(K'))^T
        k_sum = k_features.sum(
            dim=-2, keepdim=True
        )  # (batch, num_heads, 1, num_features)
        denominator = torch.matmul(
            q_features, k_sum.transpose(-2, -1)
        )  # (batch, num_heads, seq_len_q, 1)

        # 应用归一化
        context = context / (denominator + 1e-8)

        # Dropout (可选，可以加在 context 上)
        context = self.dropout(context)

        # 合并头和最终投影
        context = combine_heads(context, self.embed_dim)
        output = self.fc(context)

        return output
