import numpy as np
import torch
import torchvision
from torchvision.datasets import mnist
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim
from torch import nn
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt

# ==================== 定义超参数 ====================
train_batch_size = 64      # 训练时每个批次的样本数量
test_batch_size = 128      # 测试时每个批次的样本数量
learning_rate = 0.01       # 初始学习率
num_epoches = 20           # 训练的总轮数
momentum = 0.9             # SGD优化器的动量参数

# ==================== 数据准备和预处理 ====================
# 定义预处理函数
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL图像转换为Tensor，并自动归一化到[0,1]范围
    transforms.Normalize([0.5], [0.5])  # 标准化到[-1,1]范围，公式：(x-0.5)/0.5
])

# 下载数据，并对数据进行预处理
train_dataset = mnist.MNIST('../data/', train=True, transform=transform, download=True)
# 训练数据集：存储在../data/目录，训练模式，应用预处理变换，如果不存在则下载

test_dataset = mnist.MNIST('../data/', train=False, transform=transform)
# 测试数据集：存储在../data/目录，测试模式，应用相同的预处理变换

# 得到数据加载器
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
# 训练数据加载器：批次大小64，打乱数据顺序

test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)
# 测试数据加载器：批次大小128，不打乱数据顺序

# ==================== 可视化源数据 ====================
# 获取一批测试数据用于可视化
examples = enumerate(test_loader)  # 枚举测试数据加载器，返回索引和批次
batch_idx, (example_data, example_targets) = next(examples)  # 获取第一个批次

print(f"示例数据形状: {example_data.shape}")  # 打印数据形状，应该是 [128, 1, 28, 28]

# 显示前6个图像
fig = plt.figure(figsize=(10, 6))  # 创建图形，尺寸10x6英寸
for i in range(6):
    plt.subplot(2, 3, i+1)  # 创建2行3列的子图，当前是第i+1个
    plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
    plt.imshow(example_data[i][0], cmap='gray', interpolation='none')  # 显示灰度图像
    plt.title('真实标签: {}'.format(example_targets[i]))  # 设置标题为真实标签
    plt.xticks([])  # 移除x轴刻度
    plt.yticks([])  # 移除y轴刻度
plt.savefig('mnist_samples.png')  # 保存图像到文件
print("MNIST样本图像已保存到: mnist_samples.png")

# ==================== 构建模型 ====================
class Net(nn.Module):  # 定义神经网络类，继承自nn.Module
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Net, self).__init__()  # 调用父类构造函数
        self.flatten = nn.Flatten()  # 展平层，将多维输入一维化
        # 使用Sequential构建网络层
        self.layer1 = nn.Sequential(  # 第一个隐藏层序列
            nn.Linear(in_dim, n_hidden_1),  # 全连接层，输入维度in_dim，输出维度n_hidden_1
            nn.BatchNorm1d(n_hidden_1)  # 批归一化层，加速训练并提高稳定性
        )
        self.layer2 = nn.Sequential(  # 第二个隐藏层序列
            nn.Linear(n_hidden_1, n_hidden_2),  # 全连接层
            nn.BatchNorm1d(n_hidden_2)  # 批归一化层
        )
        self.out = nn.Sequential(  # 输出层序列
            nn.Linear(n_hidden_2, out_dim)  # 输出层，输出维度out_dim（10个类别）
        )

    def forward(self, x):  # 定义前向传播过程
        x = self.flatten(x)  # 展平输入，从[batch,1,28,28]变为[batch,784]
        x = F.relu(self.layer1(x))  # 第一层：线性变换+批归一化+ReLU激活
        x = F.relu(self.layer2(x))  # 第二层：线性变换+批归一化+ReLU激活
        x = F.softmax(self.out(x), dim=1)  # 输出层：线性变换+softmax，按行计算概率分布
        return x

# ==================== 模型实例化和配置 ====================
# 设置设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 如果有GPU则使用GPU，否则使用CPU

print(f"使用设备: {device}")

# 实例化模型
model = Net(28 * 28, 300, 100, 10)  # 输入28x28=784，隐藏层300和100，输出10个类别
model.to(device)  # 将模型移动到指定设备（GPU或CPU）

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数，适用于多分类问题
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
# SGD优化器：传入模型参数、学习率、动量

# 打印模型结构
print("模型结构:")
print(model)

# ==================== 训练模型 ====================
def main():  # 定义主函数，封装训练流程
    # 初始化记录列表
    losses = []        # 记录每个epoch的训练损失
    acces = []         # 记录每个epoch的训练准确率
    eval_losses = []   # 记录每个epoch的测试损失
    eval_acces = []    # 记录每个epoch的测试准确率
    
    # 创建TensorBoard写入器
    writer = SummaryWriter(log_dir='logs', comment='train-loss')
    # 创建SummaryWriter对象，日志目录为'logs'，注释为'train-loss'
    
    print("开始训练...")
    for epoch in range(num_epoches):  # 遍历每个训练轮次
        train_loss = 0  # 初始化当前epoch的训练损失
        train_acc = 0   # 初始化当前epoch的训练准确率
        
        # 设置为训练模式
        model.train()  # 启用dropout和batch normalization的训练模式
        
        # 动态修改参数学习率
        if epoch % 5 == 0:  # 每5个epoch调整一次学习率
            optimizer.param_groups[0]['lr'] *= 0.9  # 学习率乘以0.9（衰减）
            print("学习率: {:.6f}".format(optimizer.param_groups[0]['lr']))
        
        # 训练阶段
        for img, label in train_loader:  # 遍历训练数据加载器中的每个批次
            img = img.to(device)   # 将图像数据移动到指定设备
            label = label.to(device)  # 将标签数据移动到指定设备
            
            # 前向传播
            out = model(img)     # 模型预测输出
            loss = criterion(out, label)  # 计算损失
            
            # 反向传播
            optimizer.zero_grad()  # 清空之前的梯度
            loss.backward()        # 反向传播计算梯度
            optimizer.step()       # 更新模型参数
            
            # 记录误差
            train_loss += loss.item()  # 累加批次损失
            
            # 计算分类的准确率
            _, pred = out.max(1)  # 获取预测类别（最大概率的索引）
            num_correct = (pred == label).sum().item()  # 计算正确预测的数量
            acc = num_correct / img.shape[0]  # 计算准确率
            train_acc += acc  # 累加批次准确率
        
        # 记录训练损失和准确率
        avg_train_loss = train_loss / len(train_loader)  # 计算平均训练损失
        avg_train_acc = train_acc / len(train_loader)    # 计算平均训练准确率
        losses.append(avg_train_loss)   # 添加到训练损失列表
        acces.append(avg_train_acc)     # 添加到训练准确率列表
        
        # 保存到TensorBoard
        writer.add_scalar('Train/Loss', avg_train_loss, epoch)      # 记录训练损失
        writer.add_scalar('Train/Accuracy', avg_train_acc, epoch)   # 记录训练准确率
        
        # 在测试集上检验结果
        eval_loss = 0  # 初始化当前epoch的测试损失
        eval_acc = 0   # 初始化当前epoch的测试准确率
        
        # 将模型改为预测模式
        model.eval()  # 禁用dropout和使用训练阶段的batch normalization统计量
        
        with torch.no_grad():  # 测试阶段不需要计算梯度，节省内存和计算资源
            for img, label in test_loader:  # 遍历测试数据加载器中的每个批次
                img = img.to(device)    # 将图像数据移动到指定设备
                label = label.to(device)   # 将标签数据移动到指定设备
                
                # 前向传播
                out = model(img)      # 模型预测输出
                loss = criterion(out, label)  # 计算损失
                
                # 记录误差
                eval_loss += loss.item()  # 累加批次损失
                
                # 记录准确率
                _, pred = out.max(1)  # 获取预测类别
                num_correct = (pred == label).sum().item()  # 计算正确预测的数量
                acc = num_correct / img.shape[0]  # 计算准确率
                eval_acc += acc  # 累加批次准确率
        
        # 记录测试损失和准确率
        avg_eval_loss = eval_loss / len(test_loader)  # 计算平均测试损失
        avg_eval_acc = eval_acc / len(test_loader)    # 计算平均测试准确率
        eval_losses.append(avg_eval_loss)  # 添加到测试损失列表
        eval_acces.append(avg_eval_acc)    # 添加到测试准确率列表
        
        # 保存到TensorBoard
        writer.add_scalar('Test/Loss', avg_eval_loss, epoch)     # 记录测试损失
        writer.add_scalar('Test/Accuracy', avg_eval_acc, epoch)  # 记录测试准确率
        
        # 打印每个epoch的结果
        print('epoch: {}, Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'
              .format(epoch, avg_train_loss, avg_train_acc, avg_eval_loss, avg_eval_acc))
    
    # 关闭TensorBoard写入器
    writer.close()
    
    print('训练完成!')
    
    # ==================== 可视化训练结果 ====================
    
    # 绘制训练损失
    plt.figure(figsize=(12, 4))  # 创建图形，尺寸12x4英寸
    
    plt.subplot(1, 2, 1)  # 创建1行2列的子图，当前是第1个
    plt.title('训练损失')  # 设置子图标题
    plt.plot(np.arange(len(losses)), losses, label='Train Loss')  # 绘制训练损失曲线
    plt.plot(np.arange(len(eval_losses)), eval_losses, label='Test Loss')  # 绘制测试损失曲线
    plt.xlabel('Epoch')  # 设置x轴标签
    plt.ylabel('Loss')   # 设置y轴标签
    plt.legend()         # 显示图例
    
    plt.subplot(1, 2, 2)  # 创建1行2列的子图，当前是第2个
    plt.title('准确率')    # 设置子图标题
    plt.plot(np.arange(len(acces)), acces, label='Train Accuracy')  # 绘制训练准确率曲线
    plt.plot(np.arange(len(eval_acces)), eval_acces, label='Test Accuracy')  # 绘制测试准确率曲线
    plt.xlabel('Epoch')    # 设置x轴标签
    plt.ylabel('Accuracy') # 设置y轴标签
    plt.legend()           # 显示图例
    
    plt.tight_layout()     # 自动调整子图参数
    plt.savefig('training_results.png')  # 保存训练结果图
    print("训练结果图已保存到: training_results.png")
    
    # ==================== 最终测试 ====================
    model.eval()  # 设置为评估模式
    final_test_acc = 0   # 初始化最终测试准确数
    total_samples = 0    # 初始化总样本数
    
    with torch.no_grad():  # 不计算梯度
        for img, label in test_loader:  # 遍历测试集
            img = img.to(device)    # 移动数据到设备
            label = label.to(device)   # 移动标签到设备
            out = model(img)        # 模型预测
            _, pred = out.max(1)    # 获取预测结果
            final_test_acc += (pred == label).sum().item()  # 累加正确预测数
            total_samples += label.size(0)  # 累加总样本数
    
    final_accuracy = 100 * final_test_acc / total_samples  # 计算最终准确率百分比
    print(f'\n最终测试准确率: {final_accuracy:.2f}%')  # 打印最终准确率
    
    # ==================== 显示一些预测结果 ====================
    model.eval()  # 设置为评估模式
    with torch.no_grad():
        # 获取一批测试数据
        dataiter = iter(test_loader)  # 创建测试数据迭代器
        images, labels = next(dataiter)  # 获取下一个批次
        images, labels = images.to(device), labels.to(device)  # 移动数据到设备
        
        # 进行预测
        outputs = model(images)    # 模型预测
        _, predicted = torch.max(outputs, 1)  # 获取预测类别
        
        # 显示前12个预测结果
        fig, axes = plt.subplots(3, 4, figsize=(12, 9))  # 创建3行4列的子图
        axes = axes.ravel()  # 将子图数组展平为一维
        
        for i in range(12):  # 遍历前12个样本
            axes[i].imshow(images[i].cpu().numpy()[0], cmap='gray')  # 显示图像
            axes[i].set_title(f'预测: {predicted[i].item()}, 真实: {labels[i].item()}')  # 设置标题
            axes[i].axis('off')  # 关闭坐标轴
            
            # 如果预测错误，用红色标题突出显示
            if predicted[i] != labels[i]:
                axes[i].title.set_color('red')  # 预测错误时标题变为红色
        
        plt.tight_layout()  # 自动调整布局
        plt.savefig('prediction_results.png')  # 保存预测结果图
        print("预测结果图已保存到: prediction_results.png")

# 使用if __name__ == '__main__'保护主程序入口
if __name__ == '__main__':
    # 设置matplotlib后端为Agg，避免显示问题（不显示图形窗口，只保存到文件）
    import matplotlib
    matplotlib.use('Agg')
    
    main()  # 调用主函数开始训练