def convolve2d(input_matrix, kernel, stride=1, padding=0):
    """
    对二维矩阵进行卷积运算（互相关，不翻转核）。

    参数:
        input_matrix : list[list[float]]  输入矩阵 (H_in, W_in)
        kernel       : list[list[float]]  卷积核   (H_k,  W_k)
        stride       : int                滑动步长
        padding      : int                边缘零填充圈数

    返回:
        output : list[list[float]]  输出特征图 (H_out, W_out)
    """
    # 1. 获取尺寸
    H_in, W_in = len(input_matrix), len(input_matrix[0])
    H_k, W_k = len(kernel), len(kernel[0])

    # 2. 添加零填充（若 padding > 0）
    if padding > 0:
        # 创建填充后的新矩阵
        H_pad = H_in + 2 * padding
        W_pad = W_in + 2 * padding
        padded = [[0.0] * W_pad for _ in range(H_pad)]
        for i in range(H_in):
            for j in range(W_in):
                padded[i + padding][j + padding] = input_matrix[i][j]
    else:
        padded = input_matrix
        H_pad, W_pad = H_in, W_in

    # 3. 计算输出尺寸
    H_out = (H_pad - H_k) // stride + 1
    W_out = (W_pad - W_k) // stride + 1

    # 4. 初始化输出矩阵
    output = [[0.0] * W_out for _ in range(H_out)]

    # 5. 滑动窗口进行卷积运算
    for i in range(0, H_out):
        for j in range(0, W_out):
            # 当前窗口左上角在 padded 中的位置
            start_i = i * stride
            start_j = j * stride

            # 对窗口内元素加权求和
            s = 0.0
            for ki in range(H_k):
                for kj in range(W_k):
                    s += padded[start_i + ki][start_j + kj] * kernel[ki][kj]
            output[i][j] = s

    return output

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

# 设置设备（优先使用 GPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'使用设备: {device}')

# 定义数据预处理流程：转为张量并归一化到 [0,1]
transform = transforms.Compose([
    transforms.ToTensor(),                     # 将 PIL Image (0-255) 转为 FloatTensor (0-1)
    transforms.Normalize((0.1307,), (0.3081,)) # MNIST 全局均值与标准差，使数据零中心化
])

# 加载训练集与测试集（若本地无数据则自动下载）
train_dataset = datasets.MNIST(
    root='./data',          # 数据存放路径
    train=True,             # 训练集
    download=True,          # 若 root 下无数据则下载
    transform=transform     # 应用预处理
)

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

# 创建数据加载器
batch_size = 64
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,           # 每个 epoch 打乱顺序
    num_workers=2           # 多线程加载（Windows 下建议设为 0）
)

test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=2
)

class MNIST_CNN(nn.Module):
    """
    一个简单的卷积神经网络，包含两个卷积层和两个全连接层。
    输入尺寸: (1, 28, 28)  # 单通道灰度图
    输出: 10 个类别的 logits
    """
    def __init__(self):
        super(MNIST_CNN, self).__init__()
        # 卷积层 1: 输入 1 通道，输出 32 通道，卷积核 3x3，步长 1，填充 1（保持尺寸）
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
        # 卷积层 2: 输入 32 通道，输出 64 通道
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        # 最大池化层: 2x2 窗口，步长 2（尺寸减半）
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全连接层 1: 将卷积特征图展平后输入，尺寸计算见下文
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # 28 -> 14 -> 7（两次池化后尺寸为 7x7）
        # 全连接层 2: 输出 10 个类别分数
        self.fc2 = nn.Linear(128, 10)
        # Dropout 层，用于防止过拟合
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        # 输入 x 形状: (batch, 1, 28, 28)
        x = self.pool(F.relu(self.conv1(x)))   # 卷积1 + ReLU + 池化 -> (batch, 32, 14, 14)
        x = self.pool(F.relu(self.conv2(x)))   # 卷积2 + ReLU + 池化 -> (batch, 64, 7, 7)
        x = x.view(x.size(0), -1)              # 展平 -> (batch, 64*7*7)
        x = self.dropout(F.relu(self.fc1(x)))  # 全连接1 + ReLU + Dropout
        x = self.fc2(x)                        # 输出层 (不使用 softmax，因为 CrossEntropyLoss 自带)
        return x

model = MNIST_CNN().to(device)          # 将模型移至 GPU/CPU
criterion = nn.CrossEntropyLoss()       # 交叉熵损失（内部自动做 softmax）
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam 优化器

def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()  # 设置为训练模式（启用 Dropout、BatchNorm 等）
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()          # 清零梯度
        output = model(data)           # 前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()                # 反向传播
        optimizer.step()               # 更新参数

        running_loss += loss.item()
        _, predicted = output.max(1)   # 取 logits 最大值索引为预测类别
        total += target.size(0)
        correct += predicted.eq(target).sum().item()

        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

    train_loss = running_loss / len(train_loader)
    train_acc = 100. * correct / total
    print(f'====> Epoch {epoch} 训练集平均损失: {train_loss:.4f}, 准确率: {train_acc:.2f}%')
    return train_loss, train_acc

def test(model, device, test_loader, criterion):
    model.eval()  # 设置为评估模式（关闭 Dropout）
    test_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():  # 测试时不计算梯度，节省内存和计算
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()

    test_loss /= len(test_loader)
    test_acc = 100. * correct / total
    print(f'测试集平均损失: {test_loss:.4f}, 准确率: {test_acc:.2f}%\n')
    return test_loss, test_acc

epochs = 5
train_losses, train_accs = [], []
test_losses, test_accs = [], []

for epoch in range(1, epochs + 1):
    train_loss, train_acc = train(model, device, train_loader, optimizer, criterion, epoch)
    test_loss, test_acc = test(model, device, test_loader, criterion)

    train_losses.append(train_loss)
    train_accs.append(train_acc)
    test_losses.append(test_loss)
    test_accs.append(test_acc)

plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plt.plot(range(1, epochs+1), train_losses, label='Train Loss')
plt.plot(range(1, epochs+1), test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Loss Curve')

plt.subplot(1,2,2)
plt.plot(range(1, epochs+1), train_accs, label='Train Acc')
plt.plot(range(1, epochs+1), test_accs, label='Test Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.legend()
plt.title('Accuracy Curve')
plt.show()

# 保存模型参数（推荐方式）
torch.save(model.state_dict(), 'mnist_cnn.pth')

# 加载模型参数（用于预测）
# model = MNIST_CNN().to(device)
# model.load_state_dict(torch.load('mnist_cnn.pth'))
# model.eval()