from torch import nn 


# 定義神經網絡模型
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()  # 將 1x28x28 展平為 784
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10)  # 最終10類輸出
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

from neural_network_model import NeuralNetwork
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor



# 設置 device
device = "cuda" if torch.cuda.is_available() else "cpu"

# 下載測試數據
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)


# 從 dataloader 中取出一個 batch
data_iter = iter(test_dataloader)
images, labels = next(data_iter)

# 拿一張圖來測試
img = images[0].unsqueeze(0).to(device)  # 加 batch 維度
label = labels[0]

# 加載模型結構并加載權重
model_path = "saved_models/fashion_mnist_model.pth"
model = NeuralNetwork().to(device)
model.load_state_dict(torch.load(model_path))
model.eval() # 設置模型為測試模式


correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_dataloader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        predicted = outputs.argmax(dim=1)
        
        # 統計正確數
        correct += (predicted == labels).sum().item()
        total += labels.size(0)
        
        # 逐個打印結果（可以根據需要注釋掉）
        for i in range(len(labels)):
            print(f"真實標簽: {labels[i].item()}, 預測結果: {predicted[i].item()}")

accuracy = 100 * correct / total
print(f"\n測試集準確率: {accuracy:.2f}%")