import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import matplotlib.pyplot as plt
import numpy as np
import copy
import warnings
warnings.filterwarnings("ignore")

class Embedding(nn.Module):
    # embedding层
    def __init__(self, d_model, vocab):
        # d_model: 词嵌入的维度（转换后获得的词向量的维度）
        # vocab：词表 的大小
        super(Embedding, self).__init__()
        # 定义Embedding层
        self.lut = nn.Embedding(vocab, d_model)
        # 将参数传入类中
        self.d_model = d_model
    
    def forward(self, x):
        # x：代表输入进模型的文本通过词汇映射后的数字张量
        return self.lut(x) * math.sqrt(self.d_model)  # 陈尚d_model的开根号是为了对最后结果进行一个缩放

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        """
        位置编码器类的初始化函数
        d_model：词嵌入维度
        dropout：置0比率
        max_len：每个句子的最大长度
        """
        super(PositionalEncoding, self).__init__()
        
        # 实例化nn中预定义的Dropout层， 并将dropout传入其中，获得对象self.dropout
        self.dropout = nn.Dropout(p=dropout)
        
        # 初始化一个位置编码矩阵，他是一个0矩阵，矩阵的大小是  max_len * d_model  
        pe = torch.zeros(max_len, d_model)
        
        # 初始化一个绝对位置矩阵，在我们这里，词汇的绝对位置就是用它的索引去表示
        # 所以我们首先使用arange方法获得一个连续自然数向量，然后在使用unsqueeze方法拓展矩阵维度
        # 又因为参数传的是1， 代表矩阵拓展的位置，会使向量编程一个max_len * 1的矩阵。
        position = torch.arange(0, max_len).unsqueeze(1)
        # 绝对位置矩阵初始化之后，接下来就是考虑如何将这个位置信息加入到位置编码矩阵中
        # 最简单方法就是现将max_len * 1的绝对位置矩阵，变换成max_len * d_model形状, 然后直接覆盖原来的初始化位置编码矩阵即可。
        # 要实现这种矩阵变化，就需要一个1 * d_model形状的变化矩阵div_term，我们对这个变化矩阵的要求除了形状外，
        # 还希望他能够将自然数的绝对位置编码缩放层足够小的数字，有助于在之后的梯度下降过程中更快地收敛，这样我们就可以开始初始化
        # 得到一个自然数矩阵，但是我们这里并没有按预计的一样的初始化一个1 * d_model的矩阵
        # 只初始化了一半，即x*d_model/2的矩阵， 这是因为这里并不是真正意义的初始化了一半的矩阵
        # 我们可以把它看作初始化了两次，而两次初始化的变化矩阵不同的处理
        # 并把这两个矩阵分别填充在位置编码矩阵的偶数和奇数位置上，组成最终的位置编码矩阵。
        
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(1000.0)/ d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 这样我们就得到了位置编码矩阵pe， pe现在还只是一个二维矩阵，要想和embeddding的输出（一个三维矩阵）
        # 就必须坨镇一个维度，所以这里使用unsequeeze拓展维度
        pe = pe.unsqueeze(0)  # 将二维蟑螂扩充为三维张量
        # 最后把pe位置编码矩阵注册层模型的buffer，什么是buffer呢
        # 我们把它认为是对模型效果有帮助的，但是却不是模型结构中超参数或者参数，不需要说着优化步骤进行更新的增益
        # 注册之后我们就可以在模型保存后重新加载时盒模型结构与参数已通被加载
        self.register_buffer('pe',  pe)
        
    def forward(self, x):
        """
        forward函数的桉树是x，表示文本序列的词嵌入表示
        """
        # 在相加之前我们对pe做一些适配工作，将这个三维张量的第二位也就是句子最大长度的那一维度切片与输入的x的第二维相同即x.size(1)
        # 因为我们默认max_len为5000，一般来讲是在太大了，很难有一个句子包含5000词汇，所以要进行与输入转给你来那个的适配
        # 最后使用Variable进行封装，使其与x的样式相同，但是他是不需要进行梯度求解的，因为这里的位置信息使用的是函数计算方式，且这一步中只是数字相加，因此把requuires_grad设置成False。
        
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        # 最后就使用self.dropout对象惊喜‘丢弃’操作，并返回结果
        return self.dropout(x)

def subsequent_mask(size):
    """
    生成向后遮掩的掩码张量，参数size是掩码张量最后两个维度的大小，最后两维形成一个方阵
    """
    # 在函数中，首先定义掩码张量的形状
    attn_shape = (1, size, size)
    # 然后使用np.ones方法想这个形状中添加1元素，形成上三角阵，最后为了节约空间
    # 再使用其中的数据类型变为无符号8为整型unit8
    sub_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    # 最后将numpy类型转化为torch中的tensor，内部做一个1 - 的操作
    # 在这个其实是做了一个三角阵的反转，subsequent_mask中的每个元素都会被1减，
    # 如果是0， subsequent_mask中的该位置由0变1
    # 如果是1， subsequent_mask中的该位置由1变0
    return torch.from_numpy(1 - sub_mask)

def attention(query, key, value, mask=None, dropout=None):
    """
    注意力机制的实现，输入分别是query，key，value，mask：掩码张量，
    dropout是nn.Dropout层的实例化对象，默认为None
    """
    # 在函数中，首先去query的最后一维的大小，一般情况下就等同于我们的词嵌入维度，命名为d_k
    d_k = query.size(-1)
    # 按照注意力公式，将query与key的转置相乘，这里面key是将最后两个维度进行转置（转置后才满足矩阵乘法），key转置之后shape为[1, 4, 6]，[1, 6, 4] * [1, 4, 6] = [1, 6, 6]
    # 再除以缩放系数，就得到注意力得分张量scores
    scores = torch.matmul(query, key.transpose(-2, -1)) // math.sqrt(d_k)
    # 接着判断是否使用掩码张量
    if mask is not None:
        # 使用tensor的masked_fill方法，将掩码张量和scores张量每个位置一一比较，如果掩码张量
        # 则对应的scores张量用-1e9这个值来替换
        scores = scores.masked_fill(mask == 0, -1e9)
    # 对scores的最后一维进行softmax操作，使用F.softmax方法，第一个参数softmax对象，第二个
    # 这样获得最终的注意力张量
    p_attn = F.softmax(scores, dim=-1)
    # 之后判断是否使用dropout进行随机置零
    if dropout is not None:
        # 将p_attn传入dropout对象中进行“丢弃”处理
        p_attn = dropout(p_attn)
    #最后，根据公司将p_atten与value张量向曾获得最终的query注意力表示，同时返回注意力权重张量
    return torch.matmul(p_attn, value), p_attn

query = key = value = pe_result

attn, p_attn = attention(query, key, value)

# 首先需要定义克隆函数，因为在多头注意力机制的实现中，用到多个结构相同的线形层
# 我们将使用clone函数将他们已通初始化在同一个网络层列表对象中，之后的结构中也会用到该函数
def clone(model, N):
    # 用于生成相同网络层的克隆函数，它的参数module表示要克隆的目标网络层，N代表需要克隆的数量
    # 在函数中，我们通过for循环对module进行N次深度拷贝，使其每个module称为独立的层
    return nn.ModuleList([copy.deepcopy(model) for _ in range(N)])

# 我们使用一个类来实现多头注意力机制的处理
class MultiHeadAttention(nn.Module):
    def __init__(self, head, embedding_dim, dropout=0.1):
        """
        在类的初始化时，会传入三个参数，head代表头数，embedding_dim代表词嵌入的维度，dropout代表进行dropout操作时置零比率，默认是0.1
        """
        super(MultiHeadAttention, self).__init__()
        # 在函数中，首先使用了一个测试中常用的assert语句，判断h是否能被d_model整除
        # 这是因为我们之后要给每个头分配等量的词特征，也就是embedding_dim//head个
        assert embedding_dim % head == 0
        self.head = head  # 传入头数
        self.embedding_dim = embedding_dim
        self.dropout = nn.Dropout(p=dropout)
        self.d_k = embedding_dim // head  # 得到每个头获得的分割词向量维度d_k
        # 然后获得线形层对象，通过nn的Linear实例化，它的内部变化矩阵是embedding_dim * Embedding_dim
        self.linears = clone(nn.Linear(embedding_dim, embedding_dim), 4)
        self.attn = None
    
    def forward(self, query, key, value, mask=None):
        """前向逻辑函数，它的输入参数有4个，前三个就是注意力机制需要的Q、K、V，
        最后一个是注意力机制中可能需要的mask掩码张量，默认是None"""
        if mask is not None:  # 如果存在掩码张量
            mask = mask.unsqueeze(0)  # 使用unsqueeze拓展维度，代表多头中的第n头
        # 接着，我们获得一个batch_size的变量，它是query尺寸的第1个数字，代表有多少条样本
        batch_size = query.size(0)
        # 之后就进入多头处理环节
        # 首先利用zip将输入QKV与三个线形层组到一起，然后使用for循环，将输入QKV分别传入到线形层中
        # 完成线性变换后，开始为每个头分割输入，这里使用view方法对线性变化的结果进行维度重塑
        # 这样就意味着每个头可以获得一部分词特征组成的句子，其中的-1代表自适应维度
        # 计算机会根据这种变换自动计算这里的值，然后对第二维和第三维进行转置操作
        # 为了让代表句子长度维度和词向量维度能够相邻，这样注意力机制才能找到迟疑与句子位置的关系
        # 从attention函数中可以看到，利用的是原始输入的倒数第一和第二，这样我们就得到了每个头的
        print('----------------------------------')
        print('query-before transpose: ', query.shape)
        print('key-before transpose: ', key.shape)
        print('value-before transpose: ', value.shape)
        # 此时，query, key, value的shape为[1, 6, 4]， [batch size, sequence length, embedding dimension]
        query, key, value = \
               [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2) 
                for model, x in zip(self.linears, (query, key, value))]
        print('----------------------------------')
        print('query-after transpose: ', query.shape)
        print('key-after transpose: ', key.shape)
        print('value-after transpose: ', value.shape)
        print('----------------------------------')
        # 此时，query, key, value的shape修改为：[1, 2, 6, 2] ， [batch size, sequence length, h, embedding dimension / h]
        # 得到每个头的输入后，接下来就是将他们传入attention中
        # 这里直接调用我们之前实现的attention函数，同时也将mask和dropout传入其中
        x, self.attn = attention(query, key, value, mask, self.dropout)
        print('x- after attention: ', x.shape)  # 经过注意力机制后，输出x的shape为：[1, 2, 6, 2]， [batch size, h, sequence length, embedding dimension / h]
        # 通过多头注意力计算后，我们就得到了每个头计算结果组成的4维张量，我们需要将其转换为输入的
        # 因此这里开始进行第一步处理环节的逆操作，先对第二和第三维惊喜转置，然后使用contiguous方法
        # 这个方法的作用就是能够让转置后的张量应用view方法，否则将无法直接使用
        # 所以，下一步就是使用view方法重塑形状，变成和输入形状相同
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)  # 对x进行reshape
        print('x', x.shape)  # x的shape回到[1, 6, 4]，  [batch size, sequence length, embedding dimension]
        return self.linears[-1](x)

class PositionalwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        """
        d_model: 线形层的输入维，因为我们希望输入通过前馈全连接层后输入和输出的维度不变
        d_ff: 第二个线形层
        droupout：置零比率
        """
        super(PositionalwiseFeedForward, self).__init__()
        # 首先按照我们预期使用nn实例化了两个线形层对象
        # 他们的参数分别是d_model， d_ff和d_ff， d_model
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        # 然后我们使用nn的Dropout实例化了对象self.dropout
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x):
        """输入参数为x，代表来自上一层的输出"""
        # 首先经过第一个线性层，然后使用Funtional中relu函数进行激活,
        # 之后再使用dropout进行随机置0，最后通过第二个线性层w2，返回最终结果.
        return self.w2(self.dropout(F.relu(self.w1(x))))

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        # features：词嵌入的维度
        # eps：一个足够小的数，在规范化公式的分母中出现，房子分母为0。
        super(LayerNorm, self).__init__()
        # 根据features的形状初始化两个参数张量a2和b2，a2为元素全为1的张量，b2为元素全为0的张量。这两个张量就是规范化层的参数
        # 因为直接对上一层得到的结果做规范化公式计算，将改变结果的正常表征，因此就需要有参数作为调节因子
        # 使其既能满足规范化要求，又能不改变针对目标的表征，最后使用nn.parameter封装，代表它们是模型的参数
        self.eps = eps
        self.a2 = nn.Parameter(torch.ones(features))
        self.b2 = nn.Parameter(torch.zeros(features))
        
    def forward(self, x):
        # x来自于上一层的输出
        # 在函数中，首先对输入变量x求其最后一个维度的均值，并保持输出维度与输入维度一致
        # 接着再求最后一个维度的标准差，然后就是根据规范化公式，用x减去均值除以标准差获得规范化的结果
        # 最后对结果乘以我们的缩放茶树，即a2， *号代表同型点乘，即对应位置进行乘法操作，加上位移参数b2返回即可
        # 
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

class SubLayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        # size: 词嵌入的维度大小
        # dropout：随机置零比率
        super(SubLayerConnection, self).__init__()
        self.size = size
        self.dropout = nn.Dropout(p=dropout)
        # 实例化了规范化对象self.norm
        self.norm = LayerNorm(size)
    
    def forward(self, x, sublayer):
        # 前向逻辑函数中，接收上一层或者子层的输入作为第一个参数
        # 将该子层连接中的子层函数作为第二个函数
        # 我们首先对输出进行规范化，然后将结果传给子层处理，之后再对子层进行dropout操作
        # 随机停止一些网络中神经元的作用，来房子过拟合，最后还有一个add操作
        # 因为存在跳跃连接，所以是将输入x与dropout后的子层输出结果相加作为最终的子层连接输出
        return x + self.dropout(sublayer(self.norm(x)))

class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        # size：词嵌入维度的大小
        # self_attn：传入多头注意力子层实例化对象，并且是自注意力机制
        # feed_forward：前馈全连接层实例化对象
        # dropout：置零比率
        super(EncoderLayer, self).__init__()
        self.size = size
        # 首先将self_attn和feed_forward传入其中
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.dropout = nn.Dropout(p=dropout)
        # 编码器层中有两个子层连接结构，所以使用clones函数进行克隆
        self.sublayer = clone(SubLayerConnection(size, dropout), 2)
    
    def forward(self, x, mask):
        # x：上一层输出
        # mask：掩码张量
        # 里面就是按照结构图左侧的流程，首先通过第一个子层连接结构，其中包含多头注意力子层，然后通过第二个子层连接结构，
        #其中包含前馈全连接子层，最后返回结果
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

class Encoder(nn.Module):
    def __init__(self, layer, N):
        # layer：编码器层
        # N: 编码器层的个数
        super(Encoder, self).__init__()
        # 首先使用clones函数克隆N个编码器层放在self.layers中
        self.layers = clone(layer, N)
        # 再初始化一个规范化层，将用在编码器的最后面
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        # forward函数的输入和编码器层相同，x代表上一层的输出，mask代表掩码张量
        # 首先就是对我们克隆的编码器层进行循环，每次都会得到一个新的x
        # 这个循环的过程，就相当于输出的x经过N个编码器层的处理
        # 最后在通过规范化层的对象self.norm进行处理，最后返回结果
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout=0.1):
        # size：词嵌入的维度大小，同时也代表解码器层的尺寸
        # self_attn：多头注意力对象，也就是说这个注意力机制需要Q=K=V
        # src_attn：多头注意力对象，这里Q!=K=V
        # feed_forward：前馈全连接层对象
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.dropout = nn.Dropout(p=dropout)
        # 按照结构图使用clones函数克隆三个子层连接对象
        self.sublayer = clone(SubLayerConnection(size, dropout), 3)
    
    def forward(self, x, memory, source_mask, target_mask):
        # x：来自上一层的输出
        # mermory：来自编码器层的语义存储变量
        # 源数据掩码张量和目标数据掩码张量
        m = memory
        # 将x传入第一个子层结构，第一个子层结构分别是x和self-attn函数，因为是自注意力机制，所以Q,K,V都是x
        # 最后一个参数的目标数据掩码张量，这是要对目标数据进行遮掩，因为此时模型可能还没有生成任何目标数据
        # 比如扎起解码器准备生成第一个字符或者词汇时，我们其实已经传入了第一个字符以便计算损失
        # 但是我们不希望在生成第一个字符时模型还能利用这个信息，因此我们会将其遮掩，同样生成第二个字符或词汇时
        # 模型只能使用第一个字符或者词汇信息，第二个字符以及之后的信息都不允许模型使用
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))
        # 接着进入第二个子层，这个职称中常规的注意力机制，q是输入x；k,v是编码层输出的memory
        # 同样也传入source_mask，但是进行源数据遮掩的原因并非是抑制信息泄露，而是这笔掉对结果没有意义的字符而陈胜的注意力值
        # 以此提升模型效果和训练速度，这样就完成了第二个子层的处理
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, source_mask))
        # 最后一个子层就是前馈全连接子层，经过它的处理后，就可以返回结果，这就是我们的解码器层结构。
        return self.sublayer[2](x, self.feed_forward)

class Decoder(nn.Module):
    def __init__(self, layer, N):
        # layer：解码器层layer
        # N：解码器层的个数N
        super(Decoder, self).__init__()
        # 首先使用clones方法克隆了N个layer，然后实例化了一个规范化层
        # 因为数据走过了所有的解码器层后，最后要做规范化处理
        self.N = N
        self.layers = clone(layer, N)
        self.norm = LayerNorm(layer.size)
    
    def forward(self, x, memory, source_mask, target_mask):
        # x：数据的嵌入表示
        # memory：编码器层的输出
        # source_mask：源数据掩码张量
        # target_mask：目标数据掩码张量
        
        # 对每个层进行循环，淡然这个循环就是变量x通过每一个层的处理
        # 得出最后的结果，再进行一次规范化返回即可
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)

# nn.functional工具包装载了网络层中那些只惊喜年计算，而没有参数的层
# 将线形层和softmax计算层一起实现，因为二者的共同目标是生成最后的结构
# 因此把类的名字叫做Generator，生成器类
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        # d_model：词嵌入维度
        # vocab_size：词表大小
        super(Generator, self).__init__()
        # 首先就是使用nn中的预定义线形层进行实例化，得到一个对象self.project等待使用
        # 这个线性层的参数有两个，计时初始化函数时传进来的 d_model和vocab_size
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.project = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        # 前向逻辑函数中输入是上一层输出张量x
        # 在函数中，首先使用上一步得到的self.project对x进行线性变化
        # 然后使用F中已经实现的log_softmax进行的softmax处理
        return F.log_softmax(self.project(x), dim=-1)

# 编码器-解码器结构实现
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        # encoder：编码器对象
        # decoder：解码器对象
        # source_embed：源数据嵌入函数
        # target_embed：目标数据嵌入函数
        # generator：输出部分的类别生成器对象
        super(EncoderDecoder, self).__init__()
        # 将参数传入类中
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator
    
    def encode(self, source, source_mask):
        # 使用src_embed对source做处理，然后和source_mask一起传给self.encoder
        return self.encoder(self.src_embed(source), source_mask)
    
    def decode(self, memory, source_mask, target, target_mask):
        # 解码函数
        # 使用tgt_embed对target做处理，然后和source_mask、target_mask、memory一起传给self.decoder
        return self.decoder(self.tgt_embed(target), memory, source_mask, 
                            target_mask)
    
    def forward(self, source, target, source_mask, target_mask):
        # source：源数据
        # target：目标数据
        # source_mask、target_mask：对应的掩码张量
        
        # 在函数中，将source、source_mask传入编码函数，得到结果后，与source-mask、target、target_mask一同传给解码函数
        return self.decode(self.encode(source, source_mask), source_mask, 
                           target, target_mask)

def make_model(source_vocab, target_vocab, N=6, d_model=512,
               d_ff=2048, head=8, dropout=0.1):
    # source_vocab：源数据特征（词汇）总数
    # target_vocab：目标数据特征（词汇）总数
    # N：编码器和解码器堆叠数
    # d_model：词向量映射维度
    # d_ff：前馈全连接网络中变化矩阵的维度
    # head：多头注意力机制中多头数
    # dropout：置零比率
    
    # 首先得到一个深度拷贝命令，接下来很多结构都要进行深度拷贝，从而保证它们彼此之间相互独立，不受干扰
    c = copy.deepcopy
    # 实例化了多头注意力机制类，
    attn = MultiHeadAttention(head, d_model, dropout)
    # 然后实例化前馈全连接类
    ff = PositionalwiseFeedForward(d_model, d_ff, dropout)
    # 实例化位置编码器类
    position = PositionalEncoding(d_model, dropout)
    # 根据结构图，最外层是EncoderDecode，在EncoderDecoder中
    # 分别是编码器层，解码器层，源数据Embedding层和位置编码组成的有序结构
    # 目标数据Embedding层和位置编码组成的有序结构，以及类别生成器层
    # 在编码器层中有attention子层以及前馈全连接子层
    # 在解码器层中有两个attention子层以及前馈全连接层
    model = EncoderDecoder(
            Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
            Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
            nn.Sequential(Embedding(d_model, source_vocab), c(position)),
            nn.Sequential(Embedding(d_model, target_vocab), c(position)),
            Generator(d_model, target_vocab))
    # 模型构建完成后，接下来就是初始化模型中的参数，比如线形层的变化矩阵
    # 这里一旦判断参数的维度大于1，则会将其初始化成为一个服从均匀分布的矩阵
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model