import math
from typing import Optional, cast

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing_extensions import deprecated

from zyplib.nn.modules.base import BaseModule

from .attn_utils import combine_heads, split_heads


@deprecated('个人测试版本; 实际使用推荐 torch.nn.functional.scaled_dot_product_attention')
def my_scaled_dot_product_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    dropout: Optional[nn.Dropout] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """缩放点积注意力核心计算 (个人实现)

    Parameters
    ----------
    - `q` : `torch.Tensor`
        - 查询张量，形状: (batch_size, num_heads, seq_len_q, d_k)
    - `k` : `torch.Tensor`
        - 键张量，形状: (batch_size, num_heads, seq_len_k, d_k)
    - `v` : `torch.Tensor`
        - 值张量，形状: (batch_size, num_heads, seq_len_v, d_v)
    - `mask` : `torch.Tensor`, optional
        - 注意力掩码, by default None
    - `dropout` : `nn.Dropout`, optional
        - Dropout 层, by default None

    Returns
    ----------
    - `Tuple[torch.Tensor, torch.Tensor]`
        - (output, attention_weights)

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.attn_utils import scaled_dot_product_attention
    >>>
    >>> q = torch.randn(32, 8, 100, 64)
    >>> k = torch.randn(32, 8, 100, 64)
    >>> v = torch.randn(32, 8, 100, 64)
    >>>
    >>> output, attn_weights = scaled_dot_product_attention(q, k, v)
    >>> print(output.shape)  # torch.Size([32, 8, 100, 64])
    >>> print(attn_weights.shape)  # torch.Size([32, 8, 100, 100])
    ```
    """
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    attn_weights = F.softmax(scores, dim=-1)

    if dropout is not None:
        attn_weights = dropout(attn_weights)

    output = torch.matmul(attn_weights, v)
    return output, attn_weights


class ScaledDotProductAttention(BaseModule):
    def __init__(self, dropout: float = 0.0):
        """计算缩放点积注意力 (Scaled Dot-Product Attention)

        多头注意力机制的核心组件。

        Parameters
        ----------
        - `dropout` : `float`, optional
            - Dropout 比率, 默认为 0.0.

        """
        super().__init__()
        self.dropout = dropout

    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, num_heads, seq_len_q, d_k)`
        - `k` : `torch.Tensor`
            - 键张量, 形状: `(batch_size, num_heads, seq_len_k, d_k)`
        - `v` : `torch.Tensor`
            - 值张量, 形状: `(batch_size, num_heads, seq_len_v, d_v)`
        - `mask` : `torch.Tensor`, optional
            - 掩码张量, 形状: `(batch_size, 1, seq_len_q, seq_len_k)`, 默认为 None.

        Returns
        -------
        - **output** (`torch.Tensor`): 应用注意力后的上下文向量, 形状: `(batch_size, num_heads, seq_len_q, d_v)`.
        """
        # dropout_layer = (
        #     nn.Dropout(self.dropout) if self.dropout > 0.0 and self.training else None
        # )
        # return my_scaled_dot_product_attention(q, k, v, mask, dropout_layer)
        # my_scaled 实现的 attention 效率更低, 使用官方实现的 scaled_dot_product_attention
        # 注意：F.scaled_dot_product_attention 只返回输出，不返回注意力权重
        dp = self.dropout if self.training else 0.0
        return cast(
            tuple[torch.Tensor, torch.Tensor],
            F.scaled_dot_product_attention(q, k, v, mask, dp),
        )


@deprecated('个人测试版本; 实际使用推荐 torch.nn.MultiheadAttention')
class MultiHeadAttention(BaseModule):
    def __init__(
        self, embed_dim: int, num_heads: int, dropout: float = 0.0, bias: bool = True
    ):
        """多头注意力 (Multi-Head Attention) 模块。

        该模块实现了 "Attention Is All You Need" 中描述的多头注意力机制。
        它允许模型在不同位置共同关注来自不同表示子空间的信息。

        Parameters
        ----------
        - `embed_dim` : `int`
            - 模型的总维度 (嵌入大小)。
        - `num_heads` : `int`
            - 并行注意力头的数量。`embed_dim` 必须能被 `num_heads` 整除。
        - `dropout` : `float`, optional
            - 注意力权重的 Dropout 比率, 默认为 0.0。

        Raises
        ------
        - `ValueError`
            - 如果 `embed_dim` 不能被 `num_heads` 整除。

        Examples
        --------
        ```python
        >>> import torch
        >>> from zyplib.nn.modules.attention.multihead import MultiHeadAttention
        >>>
        >>> # 自注意力示例
        >>> embed_dim, num_heads = 128, 8
        >>> seq_len, batch_size = 10, 32
        >>> x = torch.randn(batch_size, seq_len, embed_dim)
        >>>
        >>> mha = MultiHeadAttention(embed_dim, num_heads)
        >>> output, attn_weights = mha(x, x, x)
        >>>
        >>> print(f"输入形状: {x.shape}")
        >>> print(f"输出形状: {output.shape}")
        >>> print(f"注意力权重形状: {attn_weights.shape}")
        输入形状: torch.Size([32, 10, 128])
        输出形状: torch.Size([32, 10, 128])
        注意力权重形状: torch.Size([32, 8, 10, 10])
        ```
        """
        super().__init__()
        if embed_dim % num_heads != 0:
            raise ValueError(
                f'embed_dim ({embed_dim}) 必须能被 num_heads ({num_heads}) 整除'
            )

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.d_k = embed_dim // num_heads

        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.attention = ScaledDotProductAttention(dropout=dropout)

        self.fc = nn.Linear(embed_dim, embed_dim)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        """将最后一个维度拆分为 (num_heads, d_k)。"""
        return split_heads(x, self.num_heads, self.d_k)

    def _combine_heads(self, x: torch.Tensor) -> torch.Tensor:
        """将多个头合并回原始形状。"""
        return combine_heads(x, self.embed_dim)

    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_k, embed_dim)`.
        - `v` : `torch.Tensor`
            - 值张量, 形状: `(batch_size, seq_len_v, embed_dim)`.
        - `mask` : `torch.Tensor`, optional
            - 掩码张量。其形状取决于掩码类型:
              - 填充掩码 (padding mask): `(batch_size, 1, 1, seq_len_k)`
              - 序列掩码 (sequence mask): `(batch_size, 1, seq_len_q, seq_len_k)`
              默认为 None.

        Returns
        -------
        - `tuple[torch.Tensor, torch.Tensor]`
            - 一个元组，包含:
                - **output** (`torch.Tensor`): 最终的输出张量, 形状: `(batch_size, seq_len_q, embed_dim)`.
                - **attn_weights** (`torch.Tensor`): 注意力权重, 形状: `(batch_size, num_heads, seq_len_q, seq_len_k)`.
        """
        # 1. 线性投影
        q, k, v = self.W_q(q), self.W_k(k), self.W_v(v)

        # 2. 拆分为多个头
        q = self._split_heads(q)
        k = self._split_heads(k)
        v = self._split_heads(v)

        # 3. 缩放点积注意力
        # context, attn_weights = self.attention(q, k, v, mask=mask)
        context = self.attention(q, k, v, mask=mask)

        # 4. 合并头并通过最终的线性层
        context = self._combine_heads(context)
        output = self.fc(context)

        return output


class SelfAttention(BaseModule):
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        bias: bool = True,
        **kwargs,
    ):
        """自注意力 (Self-Attention) 模块。

        MultiHeadAttention 的特化版，其中 query, key, 和 value
        都来自同一个输入序列。

        Parameters
        ----------
        - `embed_dim` : `int`
            - 模型的总维度 (嵌入大小)。
        - `num_heads` : `int`
            - 并行注意力头的数量。`embed_dim` 必须能被 `num_heads` 整除。
        - `dropout` : `float`, optional
            - 注意力权重的 Dropout 比率, 默认为 0.0。

        Examples
        --------
        ```python
        >>> import torch
        >>> from zyplib.nn.modules.attention.multihead import SelfAttention
        >>>
        >>> embed_dim, num_heads = 128, 8
        >>> seq_len, batch_size = 10, 32
        >>> x = torch.randn(batch_size, seq_len, embed_dim)
        >>>
        >>> self_attn = SelfAttention(embed_dim, num_heads)
        >>> output, attn_weights = self_attn(x)
        >>>
        >>> print(f"输入形状: {x.shape}")
        >>> print(f"输出形状: {output.shape}")
        >>> print(f"注意力权重形状: {attn_weights.shape}")
        输入形状: torch.Size([32, 10, 128])
        输出形状: torch.Size([32, 10, 128])
        注意力权重形状: torch.Size([32, 8, 10, 10])
        ```
        """
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads

        self.mha = nn.MultiheadAttention(
            embed_dim=embed_dim,
            num_heads=num_heads,
            dropout=dropout,
            bias=bias,
            batch_first=True,  # 确保输入输出形状为 [N, L, E]
            **kwargs,
        )

    def forward(
        self,
        x: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
        key_padding_mask: Optional[torch.Tensor] = None,
        need_weights=True,
        average_attn_weights=True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        自注意力的前向传播过程。

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量, 形状: `(batch_size, seq_len, embed_dim)`.
        - `attn_mask` : `torch.Tensor`, optional
            - 2D 或 3D 掩码，防止 token 关注某些其他 token。
              形状: `(seq_len, seq_len)` 或 `(batch_size * num_heads, seq_len, seq_len)`.
              布尔类型中，True 表示“不被允许关注”。
        - `key_padding_mask` : `torch.Tensor`, optional
            - 布尔掩码，指示哪些 key token 是 padding，不应被关注。
              形状: `(batch_size, seq_len)`. True 表示该位置是 padding。

        Returns
        -------
        - `tuple[torch.Tensor, torch.Tensor]`
            - **output** (`torch.Tensor`): 输出张量, 形状: `(batch_size, seq_len, embed_dim)`.
            - **attn_weights** (`torch.Tensor`): 注意力权重
                - 如果 `average_attn_weights=False`, 则返
        """
        # 对于自注意力, query, key, value 都是 x
        # need_weights=True 确保返回注意力权重
        return self.mha(
            x,
            x,
            x,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=need_weights,
            average_attn_weights=average_attn_weights,
        )


class CrossAttention(BaseModule):
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        bias: bool = True,
        **kwargs,
    ):
        """交叉注意力 (Cross-Attention) 模块。

        这是 MultiHeadAttention 的特化版本，其中 query 来自一个序列
        (例如，解码器输入)，而 key/value 来自另一个序列 (例如，编码器输出)。
        它通常用于编码器-解码器架构中。

        Parameters
        ----------
        - `embed_dim` : `int`
            - 模型的总维度 (嵌入大小)。
        - `num_heads` : `int`
            - 并行注意力头的数量。`embed_dim` 必须能被 `num_heads` 整除。
        - `dropout` : `float`, optional
            - 注意力权重的 Dropout 比率, 默认为 0.0。

        Examples
        --------
        ```python
        >>> import torch
        >>> from zyplib.nn.modules.attention.multihead import CrossAttention
        >>>
        >>> embed_dim, num_heads = 128, 8
        >>> seq_len_q, seq_len_kv, batch_size = 10, 15, 32
        >>> query = torch.randn(batch_size, seq_len_q, embed_dim)
        >>> context = torch.randn(batch_size, seq_len_kv, embed_dim)
        >>>
        >>> cross_attn = CrossAttention(embed_dim, num_heads)
        >>> output, attn_weights = cross_attn(query, context)
        >>>
        >>> print(f"Query 形状: {query.shape}")
        >>> print(f"Context 形状: {context.shape}")
        >>> print(f"输出形状: {output.shape}")
        >>> print(f"注意力权重形状: {attn_weights.shape}")
        Query 形状: torch.Size([32, 10, 128])
        Context 形状: torch.Size([32, 15, 128])
        输出形状: torch.Size([32, 10, 128])
        注意力权重形状: torch.Size([32, 8, 10, 15])
        ```
        """
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.mha = nn.MultiheadAttention(
            embed_dim=embed_dim,
            num_heads=num_heads,
            dropout=dropout,
            bias=bias,
            batch_first=True,  # 确保输入输出形状为 [N, L, E]
            **kwargs,
        )

    def forward(
        self,
        query: torch.Tensor,
        context: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
        key_padding_mask: Optional[torch.Tensor] = None,
        need_weights=True,
        average_attn_weights=True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        交叉注意力的前向传播过程。

        Parameters
        ----------
        - `query` : `torch.Tensor`
            - 查询张量, 形状: `(batch_size, seq_len_q, embed_dim)`.
        - `context` : `torch.Tensor`
            - 上下文张量 (用于 key 和 value), 形状: `(batch_size, seq_len_kv, embed_dim)`.
        - `attn_mask` : `torch.Tensor`, optional
            - 2D 掩码，防止 query token 关注某些 key token。
              形状: `(seq_len_q, seq_len_kv)`.
        - `key_padding_mask` : `torch.Tensor`, optional
            - 布尔掩码，指示哪些 key token 是 padding，不应被关注。
              形状: `(batch_size, seq_len_kv)`.

        Returns
        -------
        - `tuple[torch.Tensor, s=Fals
        """
        # 对于交叉注意力, key 和 value 都来自 context
        # average_attn_weights=False 返回每个头的注意力权重而不是平均值
        return self.mha(
            query,
            context,
            context,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=need_weights,
            average_attn_weights=average_attn_weights,
        )
