from PIL import Image, ImageFont, ImageDraw
import random
import glob
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torchvision.transforms as transforms
import numpy as np
import torch.optim as optim
import torch
import time
import copy
from create_dataset.digital_mnist_digits import PrintedMNIST
from get_model import get_model

# Define parameters
batch_size = 64
net = get_model("resnet50")
n_epochs = 10
device = "cuda" if torch.cuda.is_available() == True else "cpu"

#Define optimizer
learning_rate = 1e-3
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=learning_rate)


def load_dataset(batch_size):
    """
    Load dataset using Pytorch Dataloader
    """
    train_transform = transforms.Compose([
        transforms.RandomRotation(10),
        transforms.ToTensor(),
        # AddGaussianNoise(0, 1.0),
        # AddSPNoise(0.1),

    ])
    val_transforms = transforms.Compose([transforms.ToTensor()])

    train_set = PrintedMNIST(50000, 42, train_transform)
    val_set = PrintedMNIST(5000, 33, val_transforms)

    train_loader = DataLoader(train_set, batch_size=batch_size)
    val_loader = DataLoader(val_set, batch_size=batch_size)
    return train_loader, val_loader


def train(model, train_loader, val_loader, criterion, optimizer, num_epochs):
    # since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    train_loss_track = []
    train_acc_track = []
    val_loss_track = []
    val_acc_track = []
    for epoch in range(num_epochs):
        print(f'Epoch {epoch + 1}/{num_epochs}')
        print('-' * 10)

        # Training loop
        train_loss, train_correct = 0, 0  
        model.train()  
        for batch in train_loader:
            images, labels = batch[0].to(device), batch[1].to(device)  # load the batch to the available device (cpu/gpu)
            outputs = model(images)  
            loss = criterion(outputs, labels)
            preds = outputs.argmax(dim=1).cpu().numpy()
            optimizer.zero_grad()   
            loss.backward()  
            optimizer.step()  
            np_labels_train = labels.cpu().numpy()
            train_loss += loss.item() * batch_size  
            train_correct += np.sum(preds == np_labels_train)
        train_loss_avg = train_loss/len(train_loader.sampler)
        train_acc_avg = train_correct/len(train_loader.sampler)
        print('Train Loss: ',train_loss_avg)  
        print('Train Accuracy: ', train_acc_avg)
        train_loss_track.append(train_loss_avg)
        train_acc_track.append(train_acc_avg)

        #Validation loop
        model.eval()  
        with torch.no_grad(): 
            valid_loss, valid_correct = 0, 0 

            for batch in val_loader:
                images, labels = batch[0].to(device), batch[1].to(device)  # load the batch to the available device
                outputs = model(images)  
                loss = criterion(outputs, labels)  
                preds = outputs.argmax(dim=1).cpu().numpy()
                np_label_val = labels.cpu().numpy()
                valid_loss += loss.item() * batch_size  
                valid_correct += np.sum(preds == np_label_val)
            if valid_correct > best_acc:
                best_acc = valid_correct
                best_model_wts = copy.deepcopy(model.state_dict())
            valid_loss_avg = valid_loss/len(val_loader.sampler)
            valid_acc_avg = valid_correct/len(val_loader.sampler)
            print('Validation Loss: ', valid_loss_avg)  
            print('Validation Accuracy: ',valid_acc_avg)
            val_loss_track.append(valid_loss_avg)
            val_acc_track.append(valid_acc_avg)
            
    # Return model with best metrics          
    model.load_state_dict(best_model_wts)
    return model, train_loss_track, train_acc_track, val_loss_track, val_acc_track

if __name__ == "__main__":
    train_loader, val_loader = load_dataset(batch_size)
    model, train_loss, train_acc, val_loss, val_acc = train(net, train_loader, val_loader, criterion, optimizer, n_epochs)