如果：
    感染部位是血液，
    细菌呈革兰氏阴性，
    细菌形状为杆状，
那么：
    该细菌可能属于某一类致病菌，
    可信程度为0.8。

from pathlib import Path
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

BATCH_SIZE = 128
LEARNING_RATE = 1e-3
EPOCHS = 5

DATA_DIR = Path("data")
MODEL_PATH = Path("mnist_mlp.pt")

def get_device() -> torch.device:
    """
    Apple Silicon Mac 优先使用 MPS；
    不支持 MPS 时自动使用 CPU。
    """
    if torch.backends.mps.is_available():
        return torch.device("mps")

    return torch.device("cpu")


class MNISTMLP(nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.network = nn.Sequential(
            # 输入图片形状：
            # [batch_size, 1, 28, 28]
            #
            # 展平后：
            # [batch_size, 784]
            nn.Flatten(),

            # 输入层 -> 隐藏层
            nn.Linear(28 * 28, 128),

            # 非线性激活函数
            nn.ReLU(),

            # 隐藏层 -> 输出层
            # 10 个输出分别对应数字 0～9
            nn.Linear(128, 10),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)

def create_dataloaders() -> tuple[DataLoader, DataLoader]:
    # 将图片转换为 PyTorch 张量
    transform = transforms.ToTensor()

    training_dataset = datasets.MNIST(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=transform,
    )

    test_dataset = datasets.MNIST(
        root=DATA_DIR,
        train=False,
        download=True,
        transform=transform,
    )

    training_loader = DataLoader(
        training_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=0,
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=0,
    )

    return training_loader, test_loader

def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    loss_function: nn.Module,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
) -> tuple[float, float]:
    model.train()

    total_loss = 0.0
    correct_count = 0
    sample_count = 0

    for images, labels in dataloader:
        # 把数据放到 MPS 或 CPU
        images = images.to(device)
        labels = labels.to(device)

        # 前向传播：模型进行预测
        logits = model(images)

        # 计算预测结果与正确答案之间的误差
        loss = loss_function(logits, labels)

        # 清空上一次计算得到的梯度
        optimizer.zero_grad(set_to_none=True)

        # 反向传播：计算每个参数应该如何修改
        loss.backward()

        # 更新权重和偏置
        optimizer.step()

        batch_size = images.size(0)

        total_loss += loss.item() * batch_size
        correct_count += (
            logits.argmax(dim=1) == labels
        ).sum().item()
        sample_count += batch_size

    average_loss = total_loss / sample_count
    accuracy = correct_count / sample_count

    return average_loss, accuracy


def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    loss_function: nn.Module,
    device: torch.device,
) -> tuple[float, float]:
    model.eval()

    total_loss = 0.0
    correct_count = 0
    sample_count = 0

    # 测试阶段不需要计算梯度
    with torch.inference_mode():
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)

            logits = model(images)
            loss = loss_function(logits, labels)

            batch_size = images.size(0)

            total_loss += loss.item() * batch_size
            correct_count += (
                logits.argmax(dim=1) == labels
            ).sum().item()
            sample_count += batch_size

    average_loss = total_loss / sample_count
    accuracy = correct_count / sample_count

    return average_loss, accuracy


def main() -> None:
    torch.manual_seed(42)

    device = get_device()
    print(f"使用设备：{device}")

    training_loader, test_loader = create_dataloaders()

    model = MNISTMLP().to(device)

    # 多分类任务使用交叉熵损失
    loss_function = nn.CrossEntropyLoss()

    # Adam 根据梯度更新权重和偏置
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=LEARNING_RATE,
    )

    for epoch in range(1, EPOCHS + 1):
        training_loss, training_accuracy = train_one_epoch(
            model=model,
            dataloader=training_loader,
            loss_function=loss_function,
            optimizer=optimizer,
            device=device,
        )

        test_loss, test_accuracy = evaluate(
            model=model,
            dataloader=test_loader,
            loss_function=loss_function,
            device=device,
        )

        print(
            f"Epoch {epoch:02d}/{EPOCHS} | "
            f"训练损失：{training_loss:.4f} | "
            f"训练准确率：{training_accuracy:.2%} | "
            f"测试损失：{test_loss:.4f} | "
            f"测试准确率：{test_accuracy:.2%}"
        )

    # 只保存模型参数
    torch.save(model.state_dict(), MODEL_PATH)
    print(f"模型已保存到：{MODEL_PATH}")


if __name__ == "__main__":
    main()