# @Author：青松

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

import numpy as np
from dsx.util import print_red


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_len=512):
        """
        位置编码类，用于给输入序列添加位置信息

        参数:
            d_model: 模型的隐藏维度大小
            max_seq_len: 预期的最大序列长度

        """
        super().__init__()

        # 创建一个形状为[max_seq_len, d_model]的位置编码矩阵
        pe = torch.zeros(max_seq_len, d_model)

        # 生成位置索引序列，形状为[max_seq_len, 1]
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)

        # 计算位置编码中的分母项，用于缩放不同维度的位置信息
        # 1、torch.arange(0, d_model, 2)：生成从0到d_model-1的偶数列，例如当d_model=4时，生成[0, 2]
        # 2、使用指数还原等价计算，即e^(-log(10000.0)/d_model)，解决直接计算“大数的幂次”可能会导致 数值溢出 或 精度损失的问题
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        # 偶数列使用正弦函数计算位置编码
        pe[:, 0::2] = torch.sin(position * div_term)

        # 奇数列使用余弦函数计算位置编码
        pe[:, 1::2] = torch.cos(position * div_term)

        # 增加批次维度，形状变为[1, max_seq_len, d_model]
        pe = pe.unsqueeze(0)

        # register_buffer用于注册一个不属于模型参数但需要保存的状态张量
        # 这样可以在模型保存和加载时自动处理这个张量
        # persistent=False表示在保存模型时不将其包含在状态字典中
        self.register_buffer("pe", pe, persistent=False)

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

        参数:
            x: 输入张量，形状为[batch_size, seq_len, d_model]

        返回:
            添加了位置编码的输入张量
        """
        # 将位置编码添加到输入张量中
        # 仅使用与输入序列长度相等的位置编码部分
        x = x + self.pe[:, : x.size(1)]
        return x


class SingleHeadAttention(nn.Module):

    def __init__(self, d_model, is_print=False):
        super(SingleHeadAttention, self).__init__()
        self.hidden_size = d_model

        # Query的变换矩阵
        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        # Key的变换矩阵
        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        # Value的变换矩阵
        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)

        self.is_print = is_print

    def forward(self, q_input, k_input, v_input, attn_mask):
        if self.is_print:
            print_red("【步骤1】计算query、key和value向量 \n")
        query = self.q_proj(q_input)
        key = self.k_proj(k_input)
        value = self.v_proj(v_input)

        # 计算注意力分数：计算query和key的点积，然后除以根号hidden_size
        scores = torch.matmul(query, key.transpose(-1, -2)) / np.sqrt(self.hidden_size)

        if self.is_print:
            print_red("【步骤2】计算注意力分数")
            print("掩码前的scores:\n", scores, "\n")

        if attn_mask is not None:
            scores.masked_fill_(attn_mask == 0, -1e9)
            if self.is_print:
                print_red("【步骤 2-2】根据掩码重设无需关注的注意力分数")
                print("掩码后的scores:\n", scores, "\n")

        # 计算注意力权重：使用softmax函数，对每个样本的注意力分数进行归一化
        attention_weights = nn.functional.softmax(scores, dim=-1)

        if self.is_print:
            print_red("【步骤3】计算注意力权重")
            print("注意力权重:\n", attention_weights, "\n")

        # 计算加权和
        output = torch.matmul(attention_weights, value)

        if self.is_print:
            print_red("【步骤 4】 计算加权和")
            print("最终加权和:\n", output, "\n")

        return output, attention_weights


# 定义层归一化：LayerNorm
class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        # epsilon 是一个小值，防止除以零
        self.eps = 1e-5
        # 可学习的缩放参数（gamma）
        self.scale = nn.Parameter(torch.ones(emb_dim))
        # 可学习的平移参数（beta）
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        # 计算特征维度上的均值，用于归一化
        mean = x.mean(dim=-1, keepdim=True)
        # 计算特征维度上的方差，用于归一化（无偏=False表示使用有偏估计）
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        # 标准化输入 x
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        # 缩放和平移标准化后的输出
        return self.scale * norm_x + self.shift


class MLP(nn.Module):
    def __init__(self, d_model, feedforward_size, hidden_act="gelu"):
        super(MLP, self).__init__()

        # 升维：第一个线性层将输入维度从 emb_dim 扩展到 feedforward_size
        self.up_proj = nn.Linear(d_model, feedforward_size, bias=False)

        # 通过引入一个可学习的门控向量，模型能够决定哪些特征应该被放大（保留），哪些特征应该被抑制（过滤掉）
        self.gate_proj = nn.Linear(d_model, feedforward_size, bias=False)

        # 降维：第二个线性层将维度从 feedforward_size 压缩回 emb_dim
        self.down_proj = nn.Linear(feedforward_size, d_model, bias=False)

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

    def forward(self, x):
        # (N, T, D) -> (N, T, feedforward_size) -> (N, T, D)
        # 1、将输入映射到feedforward_size维度（升维）
        # 2、使用门控机制，实现特征选择和过滤
        # 3、使用激活函数对门控系数进行非线性变换，激活函数为gelu
        # 4、将输入映射回D维度（降维）
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


def create_cross_attn_pad_mask(query_ids, key_ids, pad_token_id):
    """
    创建填充掩码矩阵，用于屏蔽padding位置的注意力

    :param query_ids: 查询ID序列，形状为[batch_size, query_len]
    :param key_ids: 键ID序列，形状为[batch_size, key_len]
    :param pad_token_id: 填充标记ID，用于识别哪些位置是padding
    :return: 填充掩码矩阵，形状为[batch_size, query_len, key_len]，其中1表示有效位置，0表示padding位置
    """
    batch_size, query_len = query_ids.size()
    batch_size, key_len = key_ids.size()

    # 生成Key序列的padding掩码：ne(pad_token_id)返回非填充位置为True，填充位置为False（ne 是 "not equal" 的缩写，表示"不等于"操作）
    # unsqueeze(1)在第1维添加维度，形状变为[batch_size, 1, key_len]
    # byte类型（torch.uint8）只需要1个字节的存储空间，对于只包含0和1的掩码矩阵，使用byte类型可以显著减少内存占用
    key_attn_pad_mask = key_ids.data.ne(pad_token_id).unsqueeze(1).byte()

    # 将Key的填充mask信息，复制给每一个Query
    return key_attn_pad_mask.expand(batch_size, query_len, key_len)


# 生成因果注意力掩码矩阵
def create_causal_attn_mask(input_ids):
    # 生成解码时因果注意力掩码的下三角矩阵，[batch_size, tgt_len, tgt_len]
    causal_attn_shape = [input_ids.size(0), input_ids.size(1), input_ids.size(1)]
    causal_attn_mask = np.tril(np.ones(causal_attn_shape), k=0)
    causal_attn_mask = torch.from_numpy(causal_attn_mask).byte()
    return causal_attn_mask
