import csv
import torch
import mlflow

from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from train import Model

    
def get_data(IMG_SIZE: int, BATCH_SIZE: int):
    testTransformer = transforms.Compose([
        transforms.Resize(size = (IMG_SIZE, IMG_SIZE), antialias = True),
        transforms.ToTensor(),
        transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]),
    ])

    testSet = datasets.ImageFolder(root = "./test", transform = testTransformer)
    testLoader = DataLoader(testSet, batch_size = BATCH_SIZE, shuffle = False)

    return testLoader


if __name__ == '__main__':
    IMG_SIZE = 224
    BATCH_SIZE = 32
    MODEL_PATH = 'model.pth'

    mlflow.set_experiment('Skin Cancer Detection - Test set')
    mlflow.start_run()

    names = {0: 'Benign', 1: 'Malignant'}
    predictions = []
    
    test_loader = get_data(IMG_SIZE, BATCH_SIZE)
    labels = [names[i] for i in test_loader.dataset.targets]

    model = Model()
    model.load_state_dict(torch.load(MODEL_PATH))

    test_correct, test_total = 0, 0

    with torch.no_grad():
        for i, data in enumerate(test_loader):
            input, label = data
            output = model(input)
            _, predicted = torch.max(output.data, 1)
            test_total += label.size(0)
            test_correct += (predicted == label).sum().item()

            predictions.extend(predicted.tolist())

    test_accuracy = test_correct / test_total

    mlflow.log_metric('test_accuracy', test_accuracy)

    predictions = [names[pred] for pred in predictions]

    print(f'Accuracy test {test_accuracy:.2%}')

    with open('predictions.csv', 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["Predictions", "Labels"])
        
        for pred, label in zip(predictions, labels):
            writer.writerow([pred, label])