import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms
from net import FullyConnectedNet
import torch.nn.functional as F
import tqdm #让循环在运行时，自动在控制台显示一个动态更新的进度条
import os

# 指定日志目录
writer = SummaryWriter(log_dir='./mlp_mse/logs')

# 1. 判断是否使用CUDA
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')
# 2. 准备数据,Dataset 读数据,数据预处理
train_set = datasets.MNIST(root='./mlp_mse/data', 
                           train=True, download=True, transform=transforms.ToTensor())
test_set = datasets.MNIST(root='./mlp_mse/data', 
                          train=False, download=True, transform=transforms.ToTensor())
#DataLoader 负责打包、分批、迭代、打乱数据
#batch_size 每个批次多少个样本,根据显存大小设置
train_loader = DataLoader(dataset=train_set, batch_size=100, shuffle=True)
test_loader = DataLoader(dataset=test_set, batch_size=100, shuffle=False)
# 3. 创建模型
model = FullyConnectedNet()
model = model.to(device)
# 4. 确定损失函数
loss_fn = nn.MSELoss()
# 5. 创建优化器 使用梯度下降算法 参数的更新
opt = torch.optim.Adam(model.parameters())

# 早停法参数：当测试集损失连续 patience 轮不再变好时，提前终止训练
patience = 5
best_test_loss = float('inf')  # 记录到目前为止最好的测试损失
no_improve_count = 0           # 记录测试损失连续没有变好的轮数

max_epochs = 1000
os.makedirs('./mlp/model', exist_ok=True)  # 确保保存模型的目录存在
for epoch in range(max_epochs):
    # 6. 训练模型
    model.train()
    train_total_loss = 0
    for images, labels in tqdm.tqdm(train_loader, desc="train", total=len(train_loader)):
        # 将数据移动到设备
        images, labels = images.to(device), labels.to(device)
        labels = F.one_hot(labels, num_classes=10).float()
        outputs = model(images) # 前向传播
        loss = loss_fn(outputs, labels) # 计算损失

        opt.zero_grad() # 清空梯度
        loss.backward() # 反向传播 计算梯度
        opt.step() # 更新参数
        train_total_loss += loss.item()
    train_avg_loss = train_total_loss / len(train_loader)
    print(f'Epoch {epoch+1}/{max_epochs}, Loss: {train_avg_loss:.4f}')
    # 7. 测试模型
    model.eval()
    test_total_loss = 0 #一轮测试的总损失
    test_total_acc = 0 #一轮测试的总得分
    #禁用梯度计算
    #在推理时，我们不需要反向传播，因此不需要计算损失函数对参数的梯度,节省内存、加快推理速度
    with torch.inference_mode(): 
        for images, labels in tqdm.tqdm(test_loader, desc="test", total=len(test_loader)):
            images, labels = images.to(device), labels.to(device)
            labels = F.one_hot(labels, num_classes=10).float()
            outputs = model(images)
            # 计算损失
            loss = loss_fn(outputs, labels)
            test_total_loss += loss.item()

            pred = torch.argmax(outputs, dim=1)
            target = torch.argmax(labels, dim=1)
            acc = torch.eq(pred, target).float().mean()
            test_total_acc += acc.item()

    test_avg_acc = test_total_acc / len(test_loader)
    print(f'epoch:{epoch+1}, Test Accuracy: {test_avg_acc:.4f}')

    test_avg_loss = test_total_loss / len(test_loader)
    print(f'epoch:{epoch+1},Test Loss: {test_avg_loss:.4f}')
    # 记录数据
    writer.add_scalars('Loss/train', {'train_avg_loss': train_avg_loss,
    'test_avg_loss': test_avg_loss}, epoch)
    writer.add_scalar('Accuracy/test', test_avg_acc, epoch)
    # 8. 早停判断与模型保存
    if test_avg_loss < best_test_loss:
        # 测试损失变好了：更新最好成绩，重置计数器，保存当前最优模型
        best_test_loss = test_avg_loss
        no_improve_count = 0
        torch.save(model.state_dict(), './mlp/model/mnist_net_best.pth')
        print(f'测试损失改善，保存最优模型 (best loss: {best_test_loss:.4f})')
    else:
        # 测试损失没有变好：计数器加一
        no_improve_count += 1
        print(f'测试损失未改善 ({no_improve_count}/{patience})')
        if no_improve_count >= patience:
            # 连续 patience 轮没有变好，提前终止训练
            print(f'早停：测试损失连续 {patience} 轮未改善，在第 {epoch+1} 轮终止训练')
            break