"""
注意力机制模块 — Person B 负责实现

包含:
1. MultiHeadAttention: 标准多头注意力 (支持 RoPE)
2. FlashMultiHeadAttention: Flash Attention 2 加速版本

技术要点:
- Scaled Dot-Product Attention
- 支持 key_padding_mask 和 attn_mask
- Flash Attention 2 使用 torch.nn.functional.scaled_dot_product_attention
- RoPE 旋转位置编码集成
"""

from __future__ import annotations

import math
from typing import Optional

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


class MultiHeadAttention(nn.Module):
    """
    标准多头注意力机制。

    TODO [Person B]: 实现以下内容:

    __init__:
    1. Q, K, V 线性投影: nn.Linear(d_model, d_model)
    2. 输出投影: nn.Linear(d_model, d_model)
    3. Dropout

    forward(query, key, value, key_padding_mask=None, attn_mask=None):
    1. 线性投影 Q, K, V
    2. reshape 为 [B, nhead, L, d_k]
    3. (可选) 应用 RoPE 旋转位置编码
    4. 计算 attention scores: QK^T / sqrt(d_k)
    5. 应用 masks (padding mask + causal mask)
    6. Softmax + Dropout
    7. 加权求和 V
    8. reshape 回 [B, L, d_model]
    9. 输出投影
    """

    def __init__(
        self,
        d_model: int = 512,
        nhead: int = 8,
        dropout: float = 0.1,
        use_rotary_embedding: bool = False,
    ):
        super().__init__()
        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
        self.d_model = d_model
        self.nhead = nhead
        self.d_k = d_model // nhead
        self.use_rotary_embedding = use_rotary_embedding

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(p=dropout)
        self.rope: Optional[nn.Module] = None

    def forward(
        self,
        query: torch.Tensor,   # [B, L_q, D]
        key: torch.Tensor,     # [B, L_k, D]
        value: torch.Tensor,   # [B, L_v, D]
        key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, L_k]
        attn_mask: Optional[torch.Tensor] = None,             # [L_q, L_k]
        is_causal: bool = False,
    ) -> torch.Tensor:
        B, L_q, _ = query.size()
        L_k = key.size(1)
        L_v = value.size(1)

        # 1. 线性投影
        Q = self.q_proj(query)  # [B, L_q, D]
        K = self.k_proj(key)    # [B, L_k, D]
        V = self.v_proj(value)  # [B, L_v, D]

        # 2. reshape 为 [B, nhead, L, d_k]
        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]

        # 3. (可选) 应用 RoPE
        if self.use_rotary_embedding and self.rope is not None:
            Q, K = self.rope.apply_rotary_pos_emb(Q, K)

        # 4. 计算 attention scores: QK^T / sqrt(d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, H, L_q, L_k]

        # 5. 应用 masks
        if key_padding_mask is not None:
            # key_padding_mask: [B, L_k] -> [B, 1, 1, L_k]
            scores = scores.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
            )
        if is_causal:
            # 生成 causal mask
            L_q, L_k_local = scores.size(-2), scores.size(-1)
            causal_mask = torch.triu(
                torch.ones(L_q, L_k_local, device=scores.device), diagonal=1
            ).bool()
            scores = scores.masked_fill(
                causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")
            )
        if attn_mask is not None:
            # attn_mask: [L_q, L_k] -> [1, 1, L_q, L_k]
            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(0), float("-inf"))

        # 6. Softmax + Dropout
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 7. 加权求和 V
        attn_output = torch.matmul(attn_weights, V)  # [B, H, L_q, d_k]

        # 8. reshape 回 [B, L_q, d_model]
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)

        # 9. 输出投影
        output = self.out_proj(attn_output)
        return output


class FlashMultiHeadAttention(nn.Module):
    """
    Flash Attention 2 加速的多头注意力。

    TODO [Person B]: 使用 PyTorch 2.0+ 的 F.scaled_dot_product_attention 实现:
    1. 与 MultiHeadAttention 结构相同
    2. 在 forward 中使用 F.scaled_dot_product_attention(Q, K, V, attn_mask, dropout, is_causal)
    3. 会自动选择最优的 attention kernel (Flash Attention / Memory-Efficient Attention)

    注意:
    - 需要 PyTorch >= 2.0
    - is_causal=True 时自动生成因果掩码，不需要手动传入 attn_mask
    """

    def __init__(
        self,
        d_model: int = 512,
        nhead: int = 8,
        dropout: float = 0.1,
        use_rotary_embedding: bool = False,
    ):
        super().__init__()
        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
        self.d_model = d_model
        self.nhead = nhead
        self.d_k = d_model // nhead
        self.use_rotary_embedding = use_rotary_embedding
        self.dropout_p = dropout

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        self.rope: Optional[nn.Module] = None

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        key_padding_mask: Optional[torch.BoolTensor] = None,
        is_causal: bool = False,
    ) -> torch.Tensor:
        B, L_q, _ = query.size()
        L_k = key.size(1)
        L_v = value.size(1)

        # 1. 线性投影
        Q = self.q_proj(query)  # [B, L_q, D]
        K = self.k_proj(key)    # [B, L_k, D]
        V = self.v_proj(value)  # [B, L_v, D]

        # 2. reshape 为 [B, nhead, L, d_k]
        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]

        # 3. (可选) 应用 RoPE
        if self.use_rotary_embedding and self.rope is not None:
            Q, K = self.rope.apply_rotary_pos_emb(Q, K)

        # 4. 构建 attn_mask 以适配 scaled_dot_product_attention
        # PyTorch >= 2.0 支持 [B, nhead, L, d_k] 的 4D 输入
        # 注意: scaled_dot_product_attention 不允许同时设置 attn_mask 和 is_causal=True
        attn_mask: Optional[torch.Tensor] = None
        if is_causal or key_padding_mask is not None:
            attn_mask = torch.zeros(
                B, self.nhead, L_q, L_k, dtype=Q.dtype, device=Q.device
            )
            if is_causal:
                # 生成 causal mask (上三角为 -inf)
                causal_mask = torch.triu(
                    torch.ones(L_q, L_k, device=Q.device), diagonal=1
                ).bool()
                attn_mask = attn_mask.masked_fill(
                    causal_mask[None, None, :, :], float("-inf")
                )
            if key_padding_mask is not None:
                # key_padding_mask: True = padding (忽略)
                _bool_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
                _bool_mask = _bool_mask.expand(B, self.nhead, L_q, L_k)
                attn_mask = attn_mask.masked_fill(_bool_mask, float("-inf"))

        # 5. Flash Attention (PyTorch 原生)
        attn_output = F.scaled_dot_product_attention(
            Q, K, V,
            attn_mask=attn_mask,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=False,  # 已通过 attn_mask 处理
        )  # [B, H, L_q, d_k]

        # 6. reshape 回 [B, L_q, d_model]
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)

        # 7. 输出投影
        output = self.out_proj(attn_output)
        return output