from torch import nn
import math
import torch

# embedding层
class Embeddings(nn.Module):
    def __init__(self, vocab_size,d_model):
        super().__init__()
        self.lut = nn.Embedding(vocab_size,d_model)
        self.d_model = d_model
    
    def forward(self,x):
        return self.lut(x) * math.sqrt(self.d_model)
# 位置编码层
class PositionalEncoding(nn.Module):
    def __init__(self, d_model,max_len=5000,dropout=0.1) -> None:
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        pe = torch.zeros(max_len,d_model)
        
        position = torch.arange(0,max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0,d_model,2) * -math.log(1000) / d_model)
        
        pe[:,0::2] = torch.sin(position * div_term)
        pe[:,1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        self.register_buffer("pe",pe)
    
    def forward(self,x):
        x = x+ self.pe[:,:x.size(1)]
        return self.dropout(x)
#注意力机制增加掩码
def get_padding_mask(x,padding_idx):
    return (x == padding_idx).unsqueeze(1)
#注意力机制
def attention(q,k,v,mask=None,dropout=None):
    d_k = k.size(-1)
    scores = torch.matmul(q,k.transpose(-2,-1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask,-1e9)
    p_attn = torch.softmax(scores,dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn,v)
#多头注意力机制
class MultiHeadedAttention(nn.Module):
    def __init__(self, d_model,n_head,dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        #必须能被整除
        assert d_model % n_head == 0
        self.d_k = d_model // n_head
        self.n_head = n_head
        
        #
        self.W_Q = nn.Linear(d_model,d_model,bias=False)
        self.W_K = nn.Linear(d_model,d_model,bias=False)
        self.W_V = nn.Linear(d_model,d_model,bias=False)
        
        self.linear = nn.Linear(d_model,d_model,bias=False)
        #norm
        self.norm = nn.LayerNorm(d_model)
    def forward(self,q,k,v,mask=None):
        res = q
        #分头
        #shape batch_size,seq_len,d_mdoel
        q = self.W_Q(q).view(q.shape[0],q.shape[1],self.n_head,self.d_k).transpose(1,2)
        k = self.W_K(k).view(k.shape[0],k.shape[1],self.n_head,self.d_k).transpose(1,2)
        v = self.W_V(v).view(v.shape[0],v.shape[1],self.n_head,self.d_k).transpose(1,2)
        
        #掩码也需要升维
        if mask is not None:
            mask = mask.unsqueeze(1)
        #计算注意力
        context = attention(q,k,v,mask,self.dropout)
        #拼接
        concat = context.transpose(1,2).reshape(q.shape[0],-1,self.n_head * self.d_k)
        output = self.linear(concat)
        return self.norm(res+output)
#前馈神经网络
class FeedForwardNet(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        # self.relu = nn.ReLU()

        self.gelu = nn.GELU() # GELU 不会像 ReLU 那样简单地把所有负数都设置为 0，而是进行更平滑的处理
        self.norm = nn.LayerNorm(d_model) # 对每个token的 d_model 个特征进行归一化，有助于训练稳定

    def forward(self, x):
        res = x
        # x = self.relu(self.w1(x))

        x = self.gelu(self.w1(x)) # 执行顺序由内向外

        x = self.dropout(self.w2(x))
        return self.norm(x + res)
#编码器子层
class EncoderLayer(nn.Module):
    def __init__(self,d_model,n_head,d_ff,dropout=0.1):
        super().__init__()
        self.mha=MultiHeadedAttention(d_model,n_head,dropout)
        self.ffn=FeedForwardNet(d_model,d_ff,dropout)
    
    def forward(self,x,mask=None):
        output = self.mha(x,x,x,mask)
        return self.ffn(output)
#编码器层
class Encoder(nn.Module):
    def __init__(self,vocab_size,d_model,n_head,d_ff,N=6,dropout=0.1):
        super().__init__()
        self.emb = Embeddings(vocab_size,d_model)
        self.pos = PositionalEncoding(d_model=d_model,dropout=dropout)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model,n_head,d_ff,dropout)
            for _ in range(N)
        ])
        
    
    def forward(self,x,mask=None):
        x = self.emb(x)
        x = self.pos(x)
        for layer in self.layers:
            x = layer(x,mask)
        return x
#生成掩码
def get_subsequent_mask(size):
    mask_shape=(1,size,size)
    return 1-torch.tril(torch.ones(mask_shape)).byte()
#解码器层
class DecoderLayer(nn.Module):
    def __init__(self,d_model,n_head,d_ff,dropout=0.1):
        super().__init__()
        self.self_mha = MultiHeadedAttention(d_model,n_head,dropout)
        self.mha = MultiHeadedAttention(d_model,n_head,dropout)
        self.ffn = FeedForwardNet(d_model,d_ff,dropout)
        
    def forward(self,x,mask,memory,src_mask):
        x = self.self_mha(x,x,x,mask)
        x = self.mha(x,memory,memory,src_mask)
        return self.ffn(x)
#
class Decoder(nn.Module):
    def __init__(self,vocab_size,d_model,n_head,d_ff,N=6,dropout=0.1):
        super().__init__()
        self.emb = Embeddings(vocab_size,d_model)
        self.pos = PositionalEncoding(d_model=d_model,dropout=dropout)
        self.layers = nn.ModuleList([
            DecoderLayer(d_model,n_head,d_ff,dropout)
            for _ in range(N)
        ])
        
    def forward(self,x,mask,memory,src_mask):
        x =self.emb(x)
        x =self.pos(x)
        for layer in self.layers:
            x = layer(x,mask,memory,src_mask)
        return x
#生成器
class Generator(nn.Module):
    def __init__(self,d_model,vocab_size):
        super().__init__()
        self.linear = nn.Linear(d_model,vocab_size)
    def forward(self,x):
        return self.linear(x)
        #return torch.softmax(self.linear(x),dim=-1)

#模型封装
class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N=6, dropout=0.1):
        super().__init__()
        self.encoder = Encoder(src_vocab_size, d_model, n_head, d_ff, N, dropout)
        self.decoder = Decoder(tgt_vocab_size, d_model, n_head, d_ff, N, dropout)
        self.generator = Generator(d_model, tgt_vocab_size)

    def forward(self, src_x, src_mask, tgt_x, tgt_mask):
        memory = self.encoder(src_x, src_mask)
        output = self.decoder(tgt_x, tgt_mask, memory, src_mask)
        return self.generator(output)

def make_model(src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N=6, dropout=0.1):
    model = Transformer(src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N, dropout)
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

if __name__ == '__main__':
    src_vocab_size = 10
    tgt_vocab_size = 20
    d_model = 8
    n_head = 2
    d_ff = 2048
    N = 6
    dropout = 0.1

    model = make_model(src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N, dropout)
    print(model)

    # 输入数据
    src_inputs = torch.tensor([
        [1, 2, 3],
        [4, 5, 0],
    ])
    src_mask = get_padding_mask(src_inputs, 0)

    tgt_inputs = torch.tensor([
        [1, 2, 3, 4],
        [4, 5, 1, 0],
    ])
    # 处理mask
    tgt_pad_mask = get_padding_mask(tgt_inputs, 0)
    subsequent_mask = get_subsequent_mask(4)
    tgt_mask = tgt_pad_mask | subsequent_mask
    tgt_mask = tgt_mask != 0
    predict = model(src_inputs, src_mask, tgt_inputs, tgt_mask)
    print(predict.shape)