"""

 Test the final classifier. DON'T RUN UNTIL YOU ARE READY
     FOR FINAL TESTING.

"""
import torch

import torchvision
import torchvision.transforms as transforms
import cifar_challenge
import torch.nn as nn


def main():

    model = torch.load('latest_model.pt', weights_only=False)
    loss_fn = nn.CrossEntropyLoss()

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    test_set = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=128,
                                              shuffle=False, num_workers=2)

    test_loss, test_acc = cifar_challenge.evaluate(test_loader, model, loss_fn)
    
    print("\nFINAL EVALUATION RESULT...\n")
    print(f"loss: {test_loss:.4f} accuracy: {test_acc:.3f} ")

if __name__ == "__main__":
    main()