import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self, dim: int, hidden_dim:int, dropout: float):
        super().__init__()
        # 定义第一层线性变化，从输入维度到隐藏层维度
        self.fc1 = nn.Linear(dim, hidden_dim, bias = False)
        # 定义第二层线性变化，从隐藏层维度到输入维度
        self.fc2 = nn.Linear(dim, hidden_dim, bias = False)
        # 定义 dropout 层，防止过拟合
        self.dropout = nn.Dropout(dropout)
        
    def foward(self, x):
        # 前向传播
        # 首先，输入 x 经过第一层线性变换和 RELU 激活函数
        # 最后，通过最后一层线性变换和 dropout 层
        return self.dropout(self.fc2(F.relu(self.fc1(x))))

import torch
import torch.nn as nn
class LayerNorm(nn.Module):
    def __init__(self, features, eps = 1e-6):
        super.__init__()
        # 通过训练中调整 a_2 的值来更改归一化结果的缩放程度，以适应特定任务的需求
        # nn.Parameter() 允许将普通的张量转换为可学习的参数
        self.a_2 = nn.Parameter(torch.ones(features)) 
        # 通过训练调整 b_2 的值来更改归一化结果的偏移量，以便更好地适应数据分布
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self,x):
        # 统计每个样本所有维度的均值，求均值和方差
        mean = x.mean(-1,keepdim = True) # mean: [bsz, max_seq, 1]
        std = x.std(-1,keepdim = True) # std: [bsz,max_len, 1]

        return self.a_2 * (x-mean) / (std + self.eps) + self.b_2

# 注意力计算
x = x + self.attention.forward(self.attention_norm(x))
# 前馈神经网络
out = h + self.feed_froward.forward(self.fnn_norm(x))

class EncoderLayer(nn.Module):
    def __init__(self, args):
        super().__init__()
        # 每个层前面都有一个 LayerNorm, 分别在 Attention之前和 MLP之前
        self.attention_norm = LayerNorm(args.n_embd)
        # Encoder 不需要掩码，传入 is_causal = Fasle
        self.attention = MultiHeadAttention(args, is_causal = False)
        self.ffn_norm = LayerNorm(args.n_embd)
        self.ffn = MLP(args.dim, args.hidden_dim, args.dropout)
        
    def forward(self, x):
        # LayerNorm 层
        norm_x = self.attention_norm(x)
        # 注意力计算
        h = h + self.attention.forward(norm_x,norm_x,norm_x)
        # LayerNorm层
        norm_h  =self.ffn_norm(h)
        # 前馈神经网络
        out = h + self.ffn.forward(norm_h)
        return out

class Encoder(nn.Module):
    def __init__(self, args):
        super(Encoder, self).__init__()
        # 一个 Encoder 由 N 层 Encoder
        self.layers = nn.MoudleList([EncoderLayer(args) for _ in range(args.n_layer)])
        self.norm = LayerNorm(args.embd)
        
    def forward(self, x):
        # 分别通过 N 层 Encoder layer
        for layer in self.layers:
            x = layer(x)
        
        return self.norm(x)

class DecoderLayer(nn.Module):
    def __init__(self, args):
        super().__init__()
        # 需要有三个 LayerNorm, 分别在 MaskAttention, Self Attention, FFN 之前
        self.atten_norm_1 = LayerNorm(args.n_embd)
        # Decoder 的第一部分是 MaskAttention，传入 is_casual=True
        self.mask_attention = MultiHeadAttention(args, is_causal=True)
        self.atten_norm_2 = LayerNorm(args.n_embd)
        # Decoder 的第二部分是类似 Encoder 的 Attention， 传入 is_casual=False
        self.attention = MultiHeadAttention(args, is_causal = False)
        self.ffn_norm = LayerNorm(args.n_embd)
        self.ffn = MLP(args.dim, args.hidden_dim, args.dropout)
        
    def forward(self,  x, encoder_out):
        # 掩码自注意力
        norm_x = self.atten_norm_1(x)
        x = x + self.mask_attention.forward(norm_x, norm_x, norm_x)
        # 多头自注意力
        norm_x = self.atten_norm_2(x)
        h = x + self.attention.forward(norm_x, encoder_out, encoder_out)
        # 前馈神经网络
        out = h + self.ffn.forward(self.ffn_norm(h))
        return out

class Decoder(nn.Module):
    def __init__(self,args):
        super(Decoder, self).__init__()
        # 一个 Decoder 由 N 个 Decoder Layer组成
        self.layers = nn.ModuleList([DecoderLayer(args) for _ in range(args.n_layers)])
        self.norm = LayerNorm(args.n_embd)
        
    def forward(self, x, encoder_out):
        for layer in self.layers:
            x = layer(x, encoder_out)
        return self.norm(x)

[[[0],[1],[2]]] # batch_size = 1, seq_len = 3

self.token_embeddings = nn.Embedding(args.vocab_size, args.dim)

class PositionalEncoding(nn.Module):
    def __init__(self,args):
        super(PositionalEncoding,self).__init__()
        # Dropout层
        # self.dropout = nn.Dropout(args.dropout)
        
        # max_len 是序列的最大长度
        pe = torch.zeros(args.max_len, args.n_embd)
        position = torch.arrange(0, args.max_len).unsqueeze(1)
        # 计算theta
        div_term = torch.exp(torch.arrange(0, args.n_embd, 2) * - (math.log(10000.0) / args.n_embd))
        # 分别计算 sin、cos 结果
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsequeeze(0)
        self.register_buffer("pe",pe)
        
    def forward(self, x):
        # 将位置编码加到 Embedding 结果上
        x = x + self.pe[:, :x.size(1)].requires_grad_(False)
        return x

class Transformer(nn.Module):
    def __init__(self, args):
        super().__init__()
        # 必须输入词表大小和 max_len
        assert args.vocab_size is not None
        assert args.max_len is not None
        self.args = args
        self.transformer = nn.ModuleList(dict(
            wte = nn.Embedding(args.vocab_size, args.n_embd),
            wpe = PositionalEncoding(args),
            drop = nn.Dropout(args.dropout),
            encoder = Encoder(args),
            decoder = Decoder(args),
        ))
        # 最后的线性层，输入时 n_embd, 输出是词表的大小
        self.lm_head = nn.Linear(args.n_embd, args.vocab_size, bias = False)
        # 初始化所有权重
        self.apply(self._init_weights)
        
        # 查看所有参数的数量
        print("number of parameters: %.2f M" % (self.get_num_params() / 1e6, ))
        
    '''统计所有参数的数量'''
    def get_num_params(self, non_embedding = False):
        # non_embedding 是否统计 embedding 的参数
        n_params = sum(p.numel() for p in self.parameters())
        # 如果不统计 embedding 的参数，则减去即可
        if non_embedding:
            n_params -= self.transformer.wte.weight.numel()
        return n_params
    
    '''初始化权重'''
    def _init_weights(self, module):
        # 线性层和 Embedding 层初始化为正则分布
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std = 0.02)
            
    def forward(self, idx, targets = None):
        # 输入为 idx， 维度为(batch_size, seq_len, 1),
        # targets 为目标序列，用以计算loss
        device = idx.device
        batch_size, seq_len = idx.size()
        assert seq_len <= self.args.max_len, f"无法计算该序列，该序列长度为{seq_len}, 最大序列长度只有{self.args.max_len}"
        
        # 通过 self.transformer
        # 首先将输入 idx 通过 Embedding 层，得到维度为(batch_size, seq_len, n_embd)
        print("idx", idx.size())
        # 通过 Embedding 层
        tok_emb = self.transformer.wte(idx)
        print("tok_emb", tok_emb.size())
        # 然后通过位置编码
        pos_emb = self.transformer.wpe(tok_emb)
        # 进行 Dropout
        x = self.transformer.drop(pos_emb)
        # 输入 Encoder
        enc_out = self.transformer.encoder(x)
        print("enc_out", enc_out.size())
        # 输入 Decoder
        x = self.transformer.decoder(x, enc_out)
        print("x after decoder:", x.size())
        
        if targets is not None:
            # 训练阶段，如果给了 targets, 则计算 loss
            # 先通过最后的 Linear 层, 得到的维度为 (batch_size, seq_len, vocab_size)
            logits = self.lm_head(x)
            # 再跟 targets 计算交叉熵
            # logits.view(-1, logits.size(-1)) 将 logits 展平为 （batch * seq_len, vocab）
            # ignore_index = -1,这里的“-1”是填充标记，不计算损失，通常是 padding 部分
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),targets.view(-1), ignore_index = -1)
        else:
            # 推理阶段，只需要 logits, loss 为 None
            # 取 -1 是只取序列中的最后一个作为输出（即下一个要预测的字）
            logits = self.lm_head(x[:, [-1], :])
            loss = None
            
        return logits, loss

A → 预测 B  # 第 0 位 A 的输出 → 学的是预测 B
B → 预测 C  # 第 1 位 B 的输出 → 学的是预测 C
C → 预测 D  # 第 2 位 C 的输出 → 学的是预测 D