import torch
import torch.nn as nn
from transformers.activations import ACT2FN


class MultiHeadAttention(torch.nn.Module):
    """
    多头自注意力模块。

    Args:
        input_dim (int): 输入特征维度。
        output_dim (int): 输出特征维度（必须能被 num_heads 整除）。
        seq_len (int): 最大序列长度，用于生成因果掩码。
        dropout (float): Dropout 概率。
        num_heads (int): 注意力头的数量。
        qkv_bias (bool): 是否在 Q、K、V 线性投影中使用偏置。
    """

    def __init__(
        self, input_dim, output_dim, seq_len, dropout, num_heads, qkv_bias=False
    ):
        super().__init__()
        assert output_dim % num_heads == 0, "output_dim must be divisible by num_heads"

        self.output_dim = output_dim
        self.num_heads = num_heads
        self.head_dim = output_dim // num_heads  # 每个头的维度

        # 定义查询、键、值的线性层
        self.query_layer = torch.nn.Linear(
            in_features=input_dim, out_features=output_dim, bias=qkv_bias
        )
        self.key_layer = torch.nn.Linear(
            in_features=input_dim, out_features=output_dim, bias=qkv_bias
        )
        self.value_layer = torch.nn.Linear(
            in_features=input_dim, out_features=output_dim, bias=qkv_bias
        )

        # 定义输出层
        self.output_layer = torch.nn.Linear(
            in_features=output_dim, out_features=output_dim
        )

        self.dropout = torch.nn.Dropout(dropout)

        # 因果掩码：上三角矩阵，用于屏蔽未来位置的注意力
        # 注册为缓冲区，不会作为模型参数更新，但会随模型移动
        self.register_buffer(
            "mask", torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
        )

    def forward(self, x):
        """
        前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, input_dim)。

        Returns:
            torch.Tensor: 输出张量，形状为 (batch_size, seq_len, output_dim)。
        """
        batch_size, seq_len, _ = x.shape

        # 1. 通过线性层得到 Q、K、V，形状为 (batch_size, seq_len, output_dim)
        queries = self.query_layer(x)
        keys = self.key_layer(x)
        values = self.value_layer(x)

        # 2. 将 Q、K、V 拆分为多头，重塑为 (batch_size, seq_len, num_heads, head_dim)
        queries = queries.view(batch_size, seq_len, self.num_heads, self.head_dim)
        keys = keys.view(batch_size, seq_len, self.num_heads, self.head_dim)
        values = values.view(batch_size, seq_len, self.num_heads, self.head_dim)

        # 转置为 (batch_size, seq_len, num_heads, head_dim) ➡️ (batch_size, num_heads, seq_len, head_dim)
        # 这一步是为了将序列长度维度和头维度交换，方便后续的分头进行自注意力计算
        queries = queries.transpose(1, 2)
        keys = keys.transpose(1, 2)
        values = values.transpose(1, 2)

        # 3. 使用矩阵计算缩放点积注意力分数，keys 转置最后两维：(batch_size, num_heads, seq_len, head_dim) ➡️ (batch_size, num_heads, head_dim, seq_len)
        # 在PyTorch中，当使用 @ 运算符（或 torch.matmul 函数）进行矩阵乘法时，会自动处理批量维度并对最后两个维度执行标准矩阵乘法
        # 这一步计算的是每个头的注意力分数，形状为 (batch_size, num_heads, seq_len, seq_len)
        attention_scores = queries @ keys.transpose(-2, -1)

        # 4. 应用因果掩码（屏蔽未来位置）
        # 将原始掩码截断到当前序列长度，并转换为布尔类型，布尔类型张量比浮点数或整数类型占用更少的内存（每个元素 1 字节），这使得它在处理大规模数据时更加高效。
        causal_mask = self.mask.bool()[:seq_len, :seq_len]
        attention_scores.masked_fill_(causal_mask, -torch.inf)

        # 5. 计算注意力权重（softmax）并应用 dropout
        attention_weights = torch.softmax(
            attention_scores / (self.head_dim**0.5), dim=-1
        )
        attention_weights = self.dropout(attention_weights)

        # 6. 加权求和得到上下文向量，形状为 (batch_size, num_heads, seq_len, head_dim)
        context = attention_weights @ values

        # 7. 合并多头输出
        # 先转置为 (batch_size, num_tokens, num_heads, head_dim)，然后重塑为 (batch_size, num_tokens, output_dim)
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, seq_len, self.output_dim)

        # 8. 输出投影（可选）
        output = self.output_layer(context)

        return output


class LayerNorm(nn.Module):
    """
    层归一化（Layer Normalization）。

    计算公式：
        y = γ * (x - μ) / √(σ² + ε) + β

    其中：
        - x: 输入张量
        - μ: 在特征维度上计算的均值
        - σ²: 在特征维度上计算的方差（无偏估计）
        - ε: 一个小常数（默认为 1e-5），防止除零
        - γ: 可学习的缩放参数（初始化为全1）

    Args:
        normalized_shape (int): 特征维度。
    """

    def __init__(self, normalized_shape):
        super().__init__()

        # 小常数 ε，防止除零
        self.epsilon = 1e-5

        # 可学习的缩放参数 γ
        self.scale = nn.Parameter(torch.ones(normalized_shape))

        # 可学习的平移参数 β
        self.shift = nn.Parameter(torch.zeros(normalized_shape))

    def forward(self, x):
        """
        对最后一维进行归一化。

        Args:
            x (torch.Tensor): 输入张量，形状为 (..., normalized_shape)。

        Returns:
            torch.Tensor: 归一化后的张量，形状相同。
        """

        # 计算均值
        mean = x.mean(dim=-1, keepdim=True)

        # 计算方差
        var = x.var(dim=-1, keepdim=True, unbiased=False)

        # 归一化
        norm_x = (x - mean) / torch.sqrt(var + self.epsilon)

        # 缩放和平移
        return self.scale * norm_x + self.shift


class FeedForward(torch.nn.Module):
    """
    带有门控机制的前馈神经网络（GLU变体）。

    该模块包含三个线性变换：
        - up_layer: 将输入维度从 d_model 扩展到 4 * d_model
        - gate_layer: 同样将输入扩展到 4 * d_model，作为门控信号
        - down_layer: 将合并后的 4 * d_model 维度压缩回 d_model

    计算公式：
        output = down_layer(activation(gate_layer(x)) * up_layer(x) )

    其中 activation 默认为 GELU，通过门控机制实现特征选择。

    Args:
        config (dict): 配置字典，必须包含键 "d_model"（模型维度）。
        hidden_act (str): 激活函数名称，默认为 "gelu"。
    """

    def __init__(self, d_model, hidden_act="gelu"):
        super().__init__()
        feedforward_size = 4 * d_model

        # 升维线性层：将输入升维到隐藏维度
        self.up_layer = torch.nn.Linear(
            in_features=d_model, out_features=feedforward_size, bias=False
        )

        # 门控线性层：生成门控信号，用于控制信息流动
        self.gate_layer = torch.nn.Linear(
            in_features=d_model, out_features=feedforward_size, bias=False
        )

        # 降维线性层：将隐藏维度降回原始维度
        self.down_layer = torch.nn.Linear(
            in_features=feedforward_size, out_features=d_model, bias=False
        )

        # 激活函数
        self.activation_fn = ACT2FN[hidden_act]

    def forward(self, x):
        """
        前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, d_model)。

        Returns:
            torch.Tensor: 输出张量，形状与输入相同。
        """
        # 门控机制：激活(门控(x)) * 升维(x)
        gated = self.activation_fn(self.gate_layer(x)) * self.up_layer(x)

        # 降维到原始维度
        output = self.down_layer(gated)

        return output


class TransformerBlock(torch.nn.Module):
    """
    标准 Transformer 编码器块，包含多头自注意力和前馈网络，每个子层后均有残差连接和层归一化。

    Args:
        config (dict): 配置字典，必须包含以下键：
            - d_model: 模型维度
            - seq_len: 上下文长度
            - n_heads: 注意力头数
            - drop_rate: Dropout 概率
            - qkv_bias: QKV 投影是否使用偏置
    """

    def __init__(self, d_model, seq_len, n_heads, drop_rate=0.1, qkv_bias=False):
        super().__init__()

        self.attention_norm = LayerNorm(d_model)
        self.attention_layer = MultiHeadAttention(
            input_dim=d_model,
            output_dim=d_model,
            seq_len=seq_len,
            num_heads=n_heads,
            dropout=drop_rate,
            qkv_bias=qkv_bias,
        )

        self.feed_forward_norm = LayerNorm(d_model)
        self.feed_forward_layer = FeedForward(d_model)

        # Dropout 层，用于防止过拟合
        self.residual_dropout = torch.nn.Dropout(drop_rate)

    def forward(self, x):
        """
        前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, d_model)。

        Returns:
            torch.Tensor: 输出张量，形状相同。
        """
        # 注意力子层（残差连接）
        residual = x  # 保存输入作为残差连接
        x = self.attention_norm(x)  # 自注意力计算前对输入进行层归一化
        x = self.attention_layer(
            x
        )  # 应用多头注意力机制，输出形状为 [batch_size, seq_len, d_model]
        x = self.residual_dropout(x)  # 应用 Dropout 层，防止过拟合
        x = x + residual  # 将注意力子层的输出与输入残差相加，形成残差连接

        # 前馈子层（残差连接）
        residual = x  # 保存注意力子层的输出作为残差连接
        x = self.feed_forward_norm(x)  # 前馈计算前对输入进行层归一化
        x = self.feed_forward_layer(
            x
        )  # 应用前馈网络，输出形状为 [batch_size, seq_len, d_model]
        x = self.residual_dropout(x)  # 应用 Dropout 层，防止过拟合
        x = x + residual  # 将前馈子层的输出与输入残差相加，形成残差连接

        return x
