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


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

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

    def __init__(
        self, input_dim, output_dim, context_length, 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 = nn.Linear(input_dim, output_dim, bias=qkv_bias)
        self.key_layer = nn.Linear(input_dim, output_dim, bias=qkv_bias)
        self.value_layer = nn.Linear(input_dim, output_dim, bias=qkv_bias)
        self.output_layer = nn.Linear(output_dim, output_dim)  # 最后的输出投影

        self.dropout = nn.Dropout(dropout)

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

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

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

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

        # 1. 通过线性层得到 Q、K、V
        queries = self.query_layer(x)  # (batch, num_tokens, output_dim)
        keys = self.key_layer(x)
        values = self.value_layer(x)

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

        # 转置为 (batch, num_heads, num_tokens, head_dim)
        queries = queries.transpose(1, 2)
        keys = keys.transpose(1, 2)
        values = values.transpose(1, 2)

        # 3. 计算缩放点积注意力分数
        # keys 转置最后两维： (batch, num_heads, head_dim, num_tokens)
        attention_scores = queries @ keys.transpose(
            -2, -1
        )  # (batch, num_heads, num_tokens, num_tokens)

        # 4. 应用因果掩码（屏蔽未来位置）
        # 截取当前序列长度的掩码，并转换为布尔类型
        causal_mask = self.mask.bool()[:num_tokens, :num_tokens]
        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. 加权求和得到上下文向量
        context = attention_weights @ values  # (batch, num_heads, num_tokens, head_dim)

        # 7. 合并多头输出
        # 先转置为 (batch, num_tokens, num_heads, head_dim)，然后重塑为 (batch, num_tokens, output_dim)
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, num_tokens, 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(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, config, hidden_act="gelu"):
        super().__init__()
        d_model = config["d_model"]
        feedforward_size = 4 * d_model

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

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

        # 降维线性层：将隐藏维度降回原始维度
        self.down_layer = 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(nn.Module):
    """
    标准 Transformer 编码器块，包含多头自注意力和前馈网络，每个子层后均有残差连接和层归一化。

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

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

        self.attention_norm = LayerNorm(config["d_model"])
        self.attention = MultiHeadAttention(
            input_dim=config["d_model"],
            output_dim=config["d_model"],
            context_length=config["context_length"],
            num_heads=config["n_heads"],
            dropout=config["drop_rate"],
            qkv_bias=config["qkv_bias"],
        )

        self.feed_forward_norm = LayerNorm(config["d_model"])
        self.feed_forward = FeedForward(config)

        self.residual_dropout = nn.Dropout(config["drop_rate"])

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

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

        Returns:
            torch.Tensor: 输出张量，形状相同。
        """
        # 注意力子层（残差连接）
        residual = x
        x = self.attention_norm(x)
        x = self.attention(x)
        x = self.residual_dropout(x)
        x = x + residual

        # 前馈子层（残差连接）
        residual = x
        x = self.feed_forward_norm(x)
        x = self.feed_forward(x)
        x = self.residual_dropout(x)
        x = x + residual

        return x
