import torch
import numpy as np
from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data.sampler import SubsetRandomSampler
import matplotlib.pyplot as plt
import cnn
import os.path

datasetPath = './dataSet/nndataset_splitted'
# how many samples per batch to load
batch_size = 10
# percentage of training set to use as validation
test_size = 0.2
valid_size = 0.2
showBatchImages = False

# number of epochs to train the model
n_epochs = 100 # you may increase this number to train a final model

# helper function to un-normalize and display an image
def imshow(img):
    img = img / 2 + 0.5  # unnormalize
    plt.imshow(np.transpose(img, (1, 2, 0)))  # convert from Tensor image

if __name__ == '__main__':

    # convert data to a normalized torch.FloatTensor
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(20),
        transforms.Resize(size=(95,35)), #num of rows, num of colums
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

    trainData = datasets.ImageFolder(os.path.join('.',datasetPath,'train'),transform=transform)
    valData = datasets.ImageFolder(os.path.join('.',datasetPath,'val'),transform=transform)
    print(trainData.class_to_idx)
    train_loader = torch.utils.data.DataLoader(trainData, batch_size=batch_size,
        shuffle=True, num_workers=1)
    valid_loader = torch.utils.data.DataLoader(valData, batch_size=batch_size, 
        shuffle=True, num_workers=1)
    classes = [0,1]
    
    if showBatchImages:
        # obtain one batch of training images
        dataiter = iter(train_loader)
        images, labels = dataiter.next()
        images = images.numpy() # convert images to numpy for display

        # plot the images in the batch, along with the corresponding labels
        fig = plt.figure(figsize=(10, 4))
        # display 20 images
        for idx in np.arange(10):
            ax = fig.add_subplot(2, 10/2, idx+1, xticks=[], yticks=[])
            imshow(images[idx])
            ax.set_title(classes[labels[idx]])
        plt.show()

    train_on_gpu = torch.cuda.is_available()
    # create a complete CNN
    model = cnn.Net()
    print(model)

    # move tensors to GPU if CUDA is available
    if train_on_gpu:
        model.cuda()

    # specify loss function
    criterion = torch.nn.CrossEntropyLoss()

    # specify optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr = 0.003, momentum= 0.9)


    valid_loss_min = np.Inf # track change in validation loss

    for epoch in range(1, n_epochs+1):

        # keep track of training and validation loss
        train_loss = 0.0
        valid_loss = 0.0
        
        ###################
        # train the model #
        ###################
        model.train()
        for data, target in train_loader:
            # move tensors to GPU if CUDA is available
            if train_on_gpu:
                data, target = data.cuda(), target.cuda()
            # clear the gradients of all optimized variables
            optimizer.zero_grad()
            # forward pass: compute predicted outputs by passing inputs to the model
            output = model(data)
            # calculate the batch loss
            loss = criterion(output, target)
            # backward pass: compute gradient of the loss with respect to model parameters
            loss.backward()
            # perform a single optimization step (parameter update)
            optimizer.step()
            # update training loss
            train_loss += loss.item()*data.size(0)
            
        ######################    
        # validate the model #
        ######################
        model.eval()
        for data, target in valid_loader:
            # move tensors to GPU if CUDA is available
            if train_on_gpu:
                data, target = data.cuda(), target.cuda()
            # forward pass: compute predicted outputs by passing inputs to the model
            
            output = model(data)
            # calculate the batch loss
            loss = criterion(output, target)
            # update average validation loss 
            valid_loss += loss.item()*data.size(0)
        
        # calculate average losses
        train_loss = train_loss/len(train_loader.dataset)
        valid_loss = valid_loss/len(valid_loader.dataset)
            
        # print training/validation statistics 
        print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(
            epoch, train_loss, valid_loss))
        
        # save model if validation loss has decreased
        if valid_loss <= valid_loss_min:
            print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
            valid_loss_min,
            valid_loss))
            torch.save(model.state_dict(), 'nn_model.pt')
            valid_loss_min = valid_loss