import torch
import matplotlib.pyplot as plt
from torch import nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

from main import optimizer

# 设置超参数
batch_size = 64

# 定义预处理步骤
transform = transforms.Compose([
    transforms.ToTensor(),  # 转换为张量，范围 [0,1]
    transforms.Normalize((0.1307,), (0.3081,))  # 标准化：均值、方差是 MNIST 的经验值
])

# 加载训练集
train_dataset = datasets.MNIST(
    root='./data',       # 数据存放路径
    train=True,          # 训练集
    download=True,       # 如果没有就下载
    transform=transform  # 应用预处理
)

# 加载测试集
test_dataset = datasets.MNIST(
    root='./data',
    train=False,         # 测试集
    download=True,
    transform=transform
)

# 构建 DataLoader
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=batch_size,
    shuffle=True         # 打乱数据，适合训练
)

test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=batch_size,
    shuffle=False        # 测试集不需要打乱
)

# 简单测试一下
print(f"训练集大小: {len(train_dataset)}")
print(f"测试集大小: {len(test_dataset)}")

# 取一个 batch 看看形状
images, labels = next(iter(train_loader))
print(f"图片批次维度: {images.shape}")   # [batch_size, 1, 28, 28]
print(f"标签批次维度: {labels.shape}")   # [batch_size]

# 从训练集中取一个 batch
images, labels = next(iter(train_loader))
'''
# 画前 9 张图
fig, axes = plt.subplots(3, 3, figsize=(6, 6))
for i, ax in enumerate(axes.flat):
    img = images[i].squeeze().numpy()  # [1,28,28] -> [28,28]
    label = labels[i].item()
    ax.imshow(img, cmap="gray")
    ax.set_title(f"Label: {label}")
    ax.axis("off")

plt.tight_layout()
plt.show()
'''

ez = nn.Sequential(
    nn.Linear(28 * 28, 256),
    nn.ReLU(),
    nn.Linear(256, 100),
    nn.ReLU(),
    nn.Linear(100, 10),
)
criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(ez.parameters(), lr=0.002)

for images, labels in train_loader:
    images = images.view(images.size(0), -1)  # [batch_size, 28*28]
    out = ez(images)
    loss = criterion(out, labels)
    #反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"<UNK>: {loss}")

#训练结束
ez.eval()  # 关闭 dropout/batchnorm 等训练特性

correct = 0
total = 0

with torch.no_grad():  # 测试不需要计算梯度，节省显存
    for images, labels in test_loader:
        images = images.view(images.size(0), -1)  # flatten
        outputs = ez(images)                       # [batch_size, 10]

        # 取每行最大值对应的索引作为预测类别
        _, predicted = torch.max(outputs, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"测试集准确率: {correct}/{total} = {correct/total*100:.2f}%")