# ============= 标准库导入 =============
import csv      # 用于读取和写入CSV文件格式
import gzip     # 用于读取gzip压缩文件（.gz后缀）
import math     # 数学函数库，提供floor等函数
import time     # 时间相关功能，用于计算训练时间
# ============= 第三方库导入 =============
import numpy as np               # 数值计算库，提供数组操作和数学函数
import matplotlib.pyplot as plt  # 绘图库，用于可视化训练结果
# ============= PyTorch核心导入 =============
import torch         # PyTorch深度学习框架主库
from torch import nn # 神经网络模块（包含各种层、损失函数等）
from torch.utils.data import Dataset, DataLoader    # 数据集处理工具
from torch.nn.utils.rnn import pack_padded_sequence # RNN变长序列处理工具
# 继承torch.utils.data.Dataset，用于创建自定义数据集
class NameDataset(Dataset):
    # 构造函数，is_train_set=True表示训练集，False表示测试集
    def __init__(self, is_train_set=True):
        # 根据is_train_set选择不同的数据文件
        filename = 'data/names_train.csv.gz' if is_train_set else 'data/names_test.csv.gz'
        # gzip.open打开压缩文件，'rt'表示以文本模式读取
        with gzip.open(filename, 'rt') as f:
            # csv.reader创建CSV文件读取器
            reader = csv.reader(f)
            # 将CSV文件的所有行读取到列表中
            rows = list(reader)
        # 提取每行的第一个元素（名字）到names列表
        self.names = [row[0] for row in rows]
        # 记录数据集大小
        self.len = len(self.names)
        # 提取每行的第二个元素（国家）到countries列表
        self.countries = [row[1] for row in rows]
        # 获取唯一的国家列表并排序
        # set(self.countries)：获取所有不重复的国家
        # sorted()：按字母顺序排序
        # list()：转换为列表
        self.country_list = list(sorted(set(self.countries)))
        # 调用getCountryDict方法创建国家到索引的映射字典
        self.country_dict = self.getCountryDict()
        # 记录国家数量
        self.country_num = len(self.country_list)
    # 用于获取指定索引的数据
    # 参数：index - 数据索引（0到len(dataset)-1）
    # 返回：名字, 国家索引
    def __getitem__(self, index):
        # self.countries[index是数据的位置编号，不是国家的类别编号]：获取国家的名称（字符串）
        # self.country_dict["国家名"]：查询这个国家名对应的数字索引(返回的是基于国家名称得到有意义的数字编码）
        return self.names[index], self.country_dict[self.countries[index]]
    # 返回数据集大小
    def __len__(self):
        return self.len
    # 创建国家名称到索引的映射字典  {'国家名': 索引}
    def getCountryDict(self):
        # 创建空字典
        country_dict = dict()
        # enumerate(iterable, start=0)：返回索引和元素
        # idx从0开始，country_name是国家名称
        for idx, country_name in enumerate(self.country_list, 0):
            # 将国家名称映射到基于国家名称得到的有意义的数字编码
            country_dict[country_name] = idx
        return country_dict
    # 根据索引获取国家名称
    def idx2country(self, index):
        return self.country_list[index]
    # 获取国家数量
    def getCountriesNum(self):
        return self.country_num
# 计算从since开始经过的时间
# 参数：since - 起始时间（time.time()的返回值）
def time_since(since):
    s = time.time() - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)
# 继承torch.nn.Module，定义神经网络模型
class RNNClassifier(torch.nn.Module):
    # 构造函数
    def __init__(self,
                 input_size,  #   input_size: 输入维度（字符数量）
                 hidden_size, #   hidden_size: 隐藏层维度
                 output_size, #   output_size: 输出维度
                 n_layers=1,  #   n_layers: RNN层数=1
                 bidirectional=True): # bidirectional: 是否双向（True）
        # 调用父类nn.Module的构造函数
        super(RNNClassifier, self).__init__()
        # 保存隐藏层大小
        self.hidden_size = hidden_size
        # 保存RNN层数
        self.n_layers = n_layers
        # 计算方向数量：双向为2，单向为1
        self.n_directions = 2 if bidirectional else 1
        # 嵌入层：将字符索引（0-127）转换为100维向量
        # input_size: 词汇表大小（128，ASCII字符数量）
        # hidden_size: 嵌入维度（100）
        self.embedding = torch.nn.Embedding(input_size, hidden_size)
        # GRU层：门控循环单元
        # hidden_size: 输入特征维度（100）
        # hidden_size: 隐藏状态维度（100）
        # n_layers: 堆叠的GRU层数（2）
        # bidirectional: 是否双向（True）。如果是双向，输出维度为hidden_size * 2
        self.gru = torch.nn.GRU(hidden_size, hidden_size, n_layers, bidirectional=bidirectional)
        # 全连接层：将GRU输出映射到分类结果
        # hidden_size * self.n_directions: 输入特征维度（单向：100，双向：200）
        # output_size: 输出维度（国家数量）
        self.fc = torch.nn.Linear(hidden_size * self.n_directions, output_size)
    # 初始化隐藏状态 batch_size：批次大小
    def _init_hidden(self, batch_size):
        # 创建全零的隐藏状态张量
        # 形状：[n_layers * n_directions, batch_size, hidden_size]
        hidden = torch.zeros(self.n_layers * self.n_directions, batch_size, self.hidden_size)
        # 调用create_tensor函数，将张量移动到GPU（如果可用）
        return create_tensor(hidden)
    # 前向传播函数
    # input: 输入张量，形状[batch_size, seq_len]
    # seq_lengths: 每个序列的实际长度，形状[batch_size]
    def forward(self, input, seq_lengths):
        # input shape : B x S → S x B
        # .t()：转置操作，将形状从[batch_size, seq_len]变为[seq_len, batch_size]
        # GRU期望的输入形状：[seq_len, batch_size, input_size]
        input = input.t()
        # 获取批次大小
        batch_size = input.size(1)
        # 初始化隐藏状态
        hidden = self._init_hidden(batch_size)
        # 通过嵌入层：将整数索引转换为向量
        # 输入：[seq_len, batch_size]
        # 输出：[seq_len, batch_size, hidden_size]
        embedding = self.embedding(input)
        # .cpu()：将张量从GPU移动到CPU（如果它在GPU上）
        # 因为pack_padded_sequence要求lengths参数在CPU上
        seq_lengths_cpu = seq_lengths.cpu()  # 创建CPU副本
        # pack_padded_sequence：打包变长序列
        # embedding: 嵌入后的序列
        # seq_lengths_cpu: 每个序列的实际长度（必须在CPU上）
        # 作用：去除填充部分，只处理实际有效的序列长度
        # 返回值：PackedSequence对象
        gru_input = pack_padded_sequence(embedding, seq_lengths_cpu)  # 使用CPU副本
        # GRU前向传播
        # 参数：gru_input: 打包后的输入序列  hidden: 初始隐藏状态
        # 返回：output:  GRU的输出         hidden: 最后一个时间步的隐藏状态
        output, hidden = self.gru(gru_input, hidden)
        # 如果是双向GRU，需要拼接最后两个隐藏状态
        # hidden形状：[n_layers * n_directions, batch_size, hidden_size]
        # hidden[-1]: 最后一层前向的隐藏状态
        # hidden[-2]: 最后一层反向的隐藏状态
        # torch.cat(..., dim=1): 在特征维度拼接
        if self.n_directions == 2:
            hidden_cat = torch.cat((hidden[-1], hidden[-2]), dim=1)
        else:
            # 单向GRU，只取最后一层的隐藏状态
            hidden_cat = hidden[-1]
        # 全连接层：将隐藏状态映射到分类结果
        # 输出形状：[batch_size, output_size]
        fc_output = self.fc(hidden_cat)
        return fc_output
# 将名字字符串转换为ASCII码列表
def name2list(name):
    # ord(c): 返回字符c的ASCII码（整数）
    # 列表推导式：将每个字符转换为ASCII码
    arr = [ord(c) for c in name]
    # 返回：ASCII码列表, 名字长度
    return arr, len(arr)
# 将张量移动到GPU（如果可用）
def create_tensor(tensor):
    # 如果USE_GPU为True且张量不在GPU上
    if USE_GPU and not tensor.is_cuda:
        # 创建GPU设备对象
        device = torch.device("cuda:0")
        # .to(device): 将张量移动到指定设备
        tensor = tensor.to(device)
    return tensor
# 将名字和国家列表转换为PyTorch张量
# names: 名字列表，如['Smith', 'Garcia', ...]
# countries: 国家索引列表，如[0, 1, ...]
def make_tensors(names, countries):
    # 对每个名字调用name2list
    # 结果：[(ASCII列表, 长度), (ASCII列表, 长度), ...]
    sequences_and_lengths = [name2list(name)for name in names]
    # 提取ASCII列表部分
    name_sequences = [s1[0] for s1 in sequences_and_lengths]
    # 创建长度张量
    # torch.LongTensor：创建64位整数张量
    seq_lengths = torch.LongTensor([s1[1] for s1 in sequences_and_lengths])
    # 将国家列表转换为LongTensor
    countries = torch.LongTensor(countries)
    # 创建名字序列张量
    # torch.zeros: 创建全零张量
    # 形状：[batch_size, max_seq_len]
    # seq_lengths.max().item(): 获取最大序列长度（标量值）
    # dtype=torch.long: 指定为64位整数类型（嵌入层需要）
    seq_tensor = torch.zeros(len(name_sequences), seq_lengths.max().item(), dtype=torch.long)
    # enumerate(zip(...), 0): 同时遍历两个列表，idx从0开始
    # seq: ASCII列表，如[83, 109, 105, 116, 104]
    # seq_len: 序列长度，如5
    for idx, (seq, seq_len) in enumerate(zip(name_sequences, seq_lengths), 0):
        # 将ASCII列表填充到对应位置
        # seq_tensor[idx, :seq_len]: 第idx行，前seq_len列
        # torch.LongTensor(seq): 将列表转换为张量
        seq_tensor[idx, :seq_len] = torch.LongTensor(seq)
    # 按长度降序排序
    # 参数：dim=0: 在第0维（批次维度）排序  descending=True: 降序排列
    # 返回：seq_lengths: 排序后的长度     perm_idx: 排序索引（原始位置→新位置）
    seq_lengths, perm_idx = seq_lengths.sort(dim=0, descending=True)
    # 使用排序索引重新排列名字序列
    seq_tensor = seq_tensor[perm_idx]
    # 使用排序索引重新排列国家标签
    countries = countries[perm_idx]
    # 将三个张量都通过create_tensor处理（可能移动到GPU）
    return create_tensor(seq_tensor), create_tensor(seq_lengths),create_tensor(countries)

def trainModel(epoch): # 训练模型
    total_loss = 0  # 累计损失
    # enumerate(trainloader, 1): 遍历训练数据加载器，i从1开始
    # names: 一批名字，如['Smith', 'Garcia', ...]
    # countries: 对应的国家索引，如[0, 1, ...]
    for i, (names, countries) in enumerate(trainloader, 1):
        inputs, seq_lengths, target = make_tensors(names, countries)  # 将数据转换为张量
        output = classifier(inputs, seq_lengths)  # 前向传播，得到预测结果
        # 计算损失
        # criterion: CrossEntropyLoss
        # output: 模型预测，形状[batch_size, num_countries]
        # target: 真实标签，形状[batch_size]
        loss = criterion(output, target)
        optimizer.zero_grad()  # 清空梯度缓存（防止梯度累积）
        loss.backward()   # 反向传播，计算梯度
        optimizer.step()  # 更新模型参数
        total_loss += loss.item() # 累加损失值   .item(): 将单元素张量转换为Python数值
        if i % 10 == 0:  # 每10个批次打印一次
            print(f'[{time_since(start)}] Epoch {epoch} ', end='')  # 打印训练时间和当前epoch
            print(f'[{i * len(inputs)}/{len(trainset)}] ', end='')  # 打印已处理的样本数/总样本数   len(inputs): 当前批次大小
            print(f'loss={total_loss / (i * len(inputs))}')# 打印平均损失   total_loss / (i * len(inputs)): 累计损失 / 已处理样本数
    return total_loss

def testModel():  # 测试模型性能
    correct = 0  # 正确预测数
    total = len(testset)  # 测试集总样本数
    print("evaluating trained model ...")
    with torch.no_grad():  # 禁用梯度计算，节省内存和计算资源
        for i, (names, countries) in enumerate(testloader, 1):
            inputs, seq_lengths, target = make_tensors(names, countries)
            output = classifier(inputs, seq_lengths)
            # 获取预测结果
            # output.max(dim=1): 在类别维度找最大值
            # dim=1: 在第1维（类别维度）操作
            # keepdim=True: 保持维度
            # [1]: 返回最大值索引（预测的类别）
            pred = output.max(dim=1, keepdim=True)[1]
            # 统计正确预测数
            # target.view_as(pred): 将target调整为和pred相同的形状
            # pred.eq(...): 比较预测和真实值是否相等
            # .sum(): 统计True的数量
            # .item(): 转换为Python数值
            correct += pred.eq(target.view_as(pred)).sum().item()
        percent = '%.2f' % (100 * correct / total)  # 计算准确率百分比，保留两位小数
        print(f'Test set: Accuracy {correct}/{total} {percent}%')
    return correct / total  # 返回准确率（0-1之间的小数）
#Parameters
HIDDEN_SIZE = 100  # GRU隐藏层维度
BATCH_SIZE = 256   # 批次大小
N_LAYER = 2        # GRU层数
N_EPOCHS = 100     # 训练轮数
N_CHARS = 128      # 字符数量（ASCII字符0-127）
USE_GPU = True     # 是否使用GPU
# 创建训练集实例      is_train_set=False：表示这是训练集，从训练文件读取数据
trainset = NameDataset(is_train_set=True)
# 创建训练数据加载器
trainloader = DataLoader(trainset,  # 数据集
                         batch_size=BATCH_SIZE,  # 批次大小
                         shuffle=True)  # 每个epoch打乱数据顺序
# 创建测试集实例   is_train_set=False：表示这是测试集，从测试文件读取数据
testset = NameDataset(is_train_set=False)
# 创建测试数据加载器
testloader = DataLoader(testset,  # 测试数据集对象
                        batch_size=BATCH_SIZE,  # 批次大小（与训练集一致，BATCH_SIZE=256）
                        shuffle=False)  # 测试集不打乱顺序
N_COUNTRY = trainset.getCountriesNum()  # 获取国家数量

if __name__ == '__main__':  # 主程序入口（防止被其他模块导入时执行）
    classifier = RNNClassifier(N_CHARS, HIDDEN_SIZE, N_COUNTRY, N_LAYER)  # 创建RNN分类器实例
    if USE_GPU:
        device = torch.device("cuda:0")
        classifier.to(device)

    criterion = nn.CrossEntropyLoss()  # 创建交叉熵损失函数
    optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001)  # 创建Adam优化器

    start = time.time()  # 记录开始时间
    print("Training for %d epochs..." % N_EPOCHS)
    acc_list = []  # 存储每个epoch的准确率
    for epoch in range(1, N_EPOCHS + 1):
        trainModel(epoch)    # 训练
        acc = testModel()    # 测试
        acc_list.append(acc) # 记录准确率

    epoch = np.arange(1,len(acc_list) + 1, 1)  # 创建epoch编号数组：从1到len(acc_list)，步长为1
    acc_list = np.array(acc_list)  # 将准确率列表转换为numpy数组
    plt.plot(epoch,acc_list) # 绘制准确率曲线
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.grid()  # 显示网格
    plt.show()  # 显示图形窗口