"""
注意力机制模块 — 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):
    """
    标准多头注意力机制。
    """

    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)

        Q = self.q_proj(query)
        K = self.k_proj(key)
        V = self.v_proj(value)

        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)
        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)
        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)

        if self.use_rotary_embedding and self.rope is not None:
            Q, K = self.rope.apply_rotary_pos_emb(Q, K)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if key_padding_mask is not None:
            scores = scores.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
            )
        if is_causal:
            L_q_local, L_k_local = scores.size(-2), scores.size(-1)
            causal_mask = torch.triu(
                torch.ones(L_q_local, 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:
            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(0), float("-inf"))

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

        attn_output = torch.matmul(attn_weights, V)

        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)

        output = self.out_proj(attn_output)
        return output


try:
    from torch.nn.functional import scaled_dot_product_attention
    _has_flash_attn = True
except ImportError:
    _has_flash_attn = False


class FlashMultiHeadAttention(nn.Module):
    """
    Flash Attention 2 加速的多头注意力。
    
    如果 PyTorch < 2.0 或 GPU 不支持，会自动回退到标准注意力。
    """

    def __init__(
        self,
        d_model: int = 512,
        nhead: int = 8,
        dropout: float = 0.1,
        use_rotary_embedding: bool = False,
    ):
        super().__init__()
        
        if not _has_flash_attn:
            self._fallback = MultiHeadAttention(
                d_model=d_model,
                nhead=nhead,
                dropout=dropout,
                use_rotary_embedding=use_rotary_embedding,
            )
            self.d_model = d_model
            self.nhead = nhead
            self.d_k = d_model // nhead
            self.use_rotary_embedding = use_rotary_embedding
            self.rope: Optional[nn.Module] = None
            return
        
        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:
        if not _has_flash_attn:
            return self._fallback(query, key, value, key_padding_mask, None, is_causal)

        B, L_q, _ = query.size()
        L_k = key.size(1)

        Q = self.q_proj(query)
        K = self.k_proj(key)
        V = self.v_proj(value)

        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)
        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)
        V = V.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)

        if self.use_rotary_embedding and self.rope is not None:
            Q, K = self.rope.apply_rotary_pos_emb(Q, K)

        # Build an additive attention bias so we can combine causal mask and
        # key_padding_mask without conflicting with the is_causal flag.
        # F.scaled_dot_product_attention raises if both is_causal and attn_mask are set.
        if is_causal or key_padding_mask is not None:
            attn_bias = torch.zeros(B, 1, L_q, L_k, device=Q.device, dtype=Q.dtype)
            if is_causal:
                causal = torch.triu(
                    torch.ones(L_q, L_k, device=Q.device, dtype=torch.bool), diagonal=1
                )
                attn_bias = attn_bias.masked_fill(causal.unsqueeze(0).unsqueeze(0), float("-inf"))
            if key_padding_mask is not None:
                # key_padding_mask: [B, L_k] bool, True = pad
                attn_bias = attn_bias.masked_fill(
                    key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
                )
            attn_output = F.scaled_dot_product_attention(
                Q, K, V,
                attn_mask=attn_bias,
                dropout_p=self.dropout_p if self.training else 0.0,
                is_causal=False,
                scale=1.0 / math.sqrt(self.d_k),
            )
        else:
            attn_output = F.scaled_dot_product_attention(
                Q, K, V,
                attn_mask=None,
                dropout_p=self.dropout_p if self.training else 0.0,
                is_causal=False,
                scale=1.0 / math.sqrt(self.d_k),
            )

        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
        output = self.out_proj(attn_output)
        return output