import math
import os
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l

# 基于位置的前馈网络
class PositionWiseFFN(nn.Module):
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        # 第一个全连接层，将输入维度从ffn_num_input映射到ffn_num_hiddens
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        # ReLU激活函数，增加非线性特性
        self.relu = nn.ReLU()
        # 第二个全连接层，将输入维度从ffn_num_hiddens映射到ffn_num_outputs
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        # 前向传播过程
        # 通过第一个全连接层，将输入X映射到隐藏层
        # 应用ReLU激活函数，增加非线性特性
        # 通过第二个全连接层，将隐藏层映射到输出层
        return self.dense2(self.relu(self.dense1(X)))

# 使用残差连接和层归一化
class AddNorm(nn.Module):
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        # 定义一个dropout层，用于随机丢弃部分神经元
        self.dropout = nn.Dropout(dropout)
        # 定义一个层归一化层，对输入进行归一化
        self.ln = nn.LayerNorm(normalized_shape)
        
    def forward(self, X, Y):
        # 前向传播过程
        # 使用残差连接：将dropout(Y)与X相加
        # 应用层归一化：对残差进行归一化
        return self.ln(self.dropout(Y) + X)

# 实现编码器中的一个层
class EncoderBlock(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                dropout, use_bias=False, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        # 多头注意力机制
        self.attention = d2l.MultiHeadAttention(key_size, query_size,
                                               value_size, num_hiddens,
                                               num_heads, dropout, use_bias)   
        # 第一个残差连接和层归一化模块
        self.addnorm1 = AddNorm(norm_shape, dropout)
        # 位置前馈网络
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)
        # 第二个残差连接和层归一化模块
        self.addnorm2 = AddNorm(norm_shape, dropout)
        
    def forward(self, X, valid_lens):
        # 前向传播过程
        # 多头注意力机制：计算注意力权重并应用于X
        # 第一个残差连接和层归一化模块：将注意力输出与X相加并进行归一化
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        # 位置前馈网络：对第一个残差连接和层归一化模块的输出进行位置前馈网络操作
        # 第二个残差连接和层归一化模块：将位置前馈网络输出与第一个残差连接和层归一化模块的输出相加并进行归一化
        return self.addnorm2(Y, self.ffn(Y))

# Transformer编码器
class TransformerEncoder(d2l.Encoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                num_heads, num_layers, dropout, use_bias=False, **kwargs): 
        super(TransformerEncoder, self).__init__(**kwargs)
        # 词嵌入层
        self.num_hiddens = num_hiddens
        # 位置编码层
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        # 创建多个编码器块组成的序列
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        # 创建了一个空的顺序容器self.blks，用于存储多个编码器块
        self.blks = nn.Sequential()
        # 通过循环迭代的方式，逐个添加编码器块到顺序容器self.blks中
        for i in range(num_layers):
            # 使用self.blks.add_module()方法将一个新的编码器块添加到顺序容器中
            self.blks.add_module(
                "block" + str(i),
                EncoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, use_bias))
            
    def forward(self, X, valid_lens, *args):
        # 前向传播过程
        # 词嵌入层：对输入进行词嵌入操作
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))   
        # 用于存储每个编码器块中的注意力权重
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            # 编码器块：将词嵌入结果传入编码器块进行处理
            X = blk(X, valid_lens)
            # 将每个编码器块的注意力权重存储到self.attention_weights列表中的对应位置
            self.attention_weights[i] = blk.attention.attention.attention_weights
        return X

# Transformer解码器也是由多个相同的层组成
class DecoderBlock(nn.Module):
    """解码器中第 i 个块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        # 第一个多头注意力机制
        self.attention1 = d2l.MultiHeadAttention(key_size, query_size,
                                                value_size, num_hiddens,
                                                num_heads, dropout)
        # 第一个残差连接和层归一化模块
        self.addnorm1 = AddNorm(norm_shape, dropout)
        # 第二个多头注意力机制
        self.attention2 = d2l.MultiHeadAttention(key_size, query_size,
                                                value_size, num_hiddens,
                                                num_heads, dropout)
        # 第二个残差连接和层归一化模块
        self.addnorm2 = AddNorm(norm_shape, dropout)
        # 位置前馈网络
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)   
        # 第三个残差连接和层归一化模块
        self.addnorm3 = AddNorm(norm_shape, dropout)
        
    def forward(self, X, state):
        # 从state中提取编码器输出的相关信息
        enc_outputs, enc_valid_lens = state[0], state[1]
        # 检查状态中当前解码器块的键-值对是否为None
        if state[2][self.i] is None:
            key_values = X
        else:
            # 将当前解码器块的键-值对与输入X在维度1上进行拼接
            key_values = torch.cat((state[2][self.i], X), axis=1)
        # 更新状态中当前解码器块的键-值对
        state[2][self.i] = key_values
        if self.training:
            batch_size, num_steps, _ = X.shape
            # 生成一个从1到num_steps的张量，用于表示解码器输入的有效长度
            dec_valid_lens = torch.arange(1, num_steps + 1,
                                         device = X.device).repeat(batch_size, 1)   
        else:
            # 在评估模式下，解码器输入的有效长度为None
            dec_valid_lens = None
        # 第一个多头注意力机制的输出
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        # 第一个残差连接和层归一化模块的输出
        Y = self.addnorm1(X, X2)
        # 第二个多头注意力机制的输出
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        # 第二个残差连接和层归一化模块的输出
        Z = self.addnorm2(Y, Y2)
        # 最终输出经过第三个残差连接和层归一化模块处理，返回结果和更新后的状态
        return self.addnorm3(Z, self.ffn(Z)), state

# Transform解码器
class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        # 初始化参数
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        # 词嵌入层
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        # 位置编码层
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        # 创建多个解码器块组成的序列
        self.blks = nn.Sequential()
        # 通过循环迭代的方式，逐个添加解码器块到顺序容器self.blks中
        for i in range(num_layers):
            # 使用self.blks.add_module()方法将一个新的解码器块添加到顺序容器中
            self.blks.add_module(
                "block" + str(i),
                DecoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, i))
        # 全连接层，将解码器块的输出转换为词汇表大小的输出
        self.dense = nn.Linear(num_hiddens, vocab_size)
        
    def init_state(self, enc_outputs, enc_valid_lens, *args):
        # 初始化解码器状态
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]
    
    def forward(self, X, state):
        # 前向传播过程
        # 词嵌入层：对输入进行词嵌入操作
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))  
        # 用于存储每个解码器块中的注意力权重
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)] 
        for i, blk in enumerate(self.blks):
            # 解码器块：将词嵌入结果传入解码器块进行处理
            X, state = blk(X, state)
            # 将每个解码器块的注意力权重存储到self._attention_weights列表中的对应位置
            self._attention_weights[0][i] = blk.attention1.attention.attention_weights   
            self._attention_weights[1][i] = blk.attention2.attention.attention_weights
        # 全连接层：将解码器块的输出转换为词汇表大小的输出
        return self.dense(X), state
    
    @property
    def attention_weights(self):
        # 返回注意力权重
        return self._attention_weights

def read_data_nmt():
    """载入 “英语-法语” 数据集 """
    # 下载并解压数据集
    data_dir = d2l.download_extract('fra-eng')
    # 打开文件并读取数据
    with open(os.path.join(data_dir, 'fra.txt'), 'r', encoding='utf-8') as f:
        return f.read()

def preprocess_nmt(text):
    """预处理 “英语-法语” 数据集"""
    def no_space(char, prev_char):
        return char in set(',.!?') and prev_char != ''
    # 替换特殊字符并转换为小写
    text = text.replace('\u202f', ' ').replace('\xa0',' ').lower()
    # 在标点符号前添加空格，以便于分词
    out = [
        ' ' + char if i > 0 and no_space(char, text[i - 1]) else char
        for i, char in enumerate(text)]
    # 将字符列表拼接为字符串
    return ''.join(out)

def tokenize_nmt(text, num_examples=None):
    """词元化 “英语-法语” 数据数据集 """
    source, target = [], []
    # 按行遍历文本
    for i, line in enumerate(text.split('\n')):
        # 检查是否达到指定的样本数量
        if num_examples and i > num_examples:
            break
        # 按制表符分割行，将源语言和目标语言分别存储到source和target列表中
        parts = line.split('\t')
        if len(parts) == 2:
            source.append(parts[0].split(' '))
            target.append(parts[1].split(' '))
    return source, target

def truncate_pad(line, num_steps, padding_token):
    """截断或填充文本序列"""
    # 检查文本序列长度是否超过指定的最大长度
    if len(line) > num_steps:
        # 如果超过最大长度，则截断序列，只保留前面的num_steps个元素
        return line[:num_steps]
    # 如果未超过最大长度，则使用padding_token进行填充，使序列长度达到num_steps
    return line + [padding_token] * (num_steps - len(line))

def build_array_nmt(lines, vocab, num_steps):
    """将机器翻译的文本序列转换成小批量"""
    lines = [vocab[l] for l in lines]
    lines = [l + [vocab['<eos>']] for l in lines]
    array = torch.tensor([ truncate_pad(l, num_steps, vocab['<pad>']) for l in lines ])
    valid_len = (array != vocab['<pad>']).type(torch.int32).sum(1)
    return array, valid_len

def load_data_nmt(batch_size, num_steps, num_examples=600):
    """返回翻译数据集的迭代器和词汇表"""
    # 载入并预处理文本数据
    text = preprocess_nmt(read_data_nmt())
    # 对文本数据进行词元化
    source, target = tokenize_nmt(text, num_examples)
    # 构建源语言和目标语言的词汇表
    src_vocab = d2l.Vocab(source, min_freq=2,
                         reserved_tokens=['<pad>','<bos>','<eos>'])
    tgt_vocab = d2l.Vocab(target, min_freq=2,
                         reserved_tokens=['<pad>','<bos>','<eos>'])
    # 将词元化后的文本数据转换为数值数组，并计算有效长度
    src_array, src_valid_len = build_array_nmt(source, src_vocab, num_steps)  
    tgt_array, tgt_valid_len = build_array_nmt(target, tgt_vocab, num_steps)
    # 构建数据集迭代器
    data_arrays = (src_array, src_valid_len, tgt_array, tgt_valid_len)
    data_iter = d2l.load_array(data_arrays, batch_size)
    # 返回数据集迭代器和词汇表
    return data_iter, src_vocab, tgt_vocab