#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Nov 12 23:43:31 2018

@author: myidispg
"""

# import libraries
import torch
import numpy as np

from torchvision import datasets
import torchvision.transforms as transforms

# Number of subprocesses to use for data loading
num_workers = 0
# how many samples to load per batch
batch_size = 20

# convert data to torch.FloatTensor()
transform = transforms.ToTensor()

# choose the training and test dataset

train_data = datasets.MNIST(root='MNIST_data', train=True, download=True, transform=transform)
test_data = datasets.MNIST(root='MNIST_data', train=True, download=True, transform=transform)

# Prepare the dataloaders
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, num_workers = num_workers)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, num_workers=num_workers)

import matplotlib.pyplot as plt
%matplotlib inline

# obtain one batch of training images
dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy()

# plot the images in the batch along with the corresponding labels
fig = plt.figure(figsize=(25,4))
for idx in np.arange(20):
    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
    ax.imshow(np.squeeze(images[idx]), cmap='gray')
    # print out the correct label for each image
    # .item() gets the value contained in a Tensor
    ax.set_title(str(labels[idx].item()))
    
# View an image in more detail
img  = np.squeeze(images[1])
fig = plt.figure(figsize=(12,12))
ax = fig.add_subplot(111)
ax.imshow(img, cmap='gray')
width, height = img.shape
thresh = img.max()/2.5
for x in range(width):
    for y in range(height):
        val = round(img[x][y], 2) if img[x][y] != 0 else 0
        ax.annotate(str(val), xy=(y,x),
                    horizontalalignment='center',
                    verticalalignment='center',
                    color='white' if img[x][y]<thresh else 'black')
        
# DEFINE THE NETOWRK ARCHITECTURE
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class Net(nn.Module):
    def __init__(self, input_size=784, hidden_1=128, hidden_2 = 64, output = 10):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_1)
        self.fc2 = nn.Linear(hidden_1, hidden_2)
        self.fc3 = nn.Linear(hidden_2, output)
        
        self.dropout = nn.Dropout(0.3)
    
    def forward(self, x):
        # flatten image input
        x = x.view(-1, 28*28)        
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.dropout(F.relu(self.fc2(x)))
        x = self.fc3(x)
        
        return x
    
# initialize the nn
model = Net()
# specify the loss function
criterion = nn.CrossEntropyLoss()
# specify optimizer
optimizer = optim.SGD(model.parameters(), lr=0.03)

# number of epochs
epochs = 30

model.train() # set the model to train mode

for epoch in range(epochs):
    train_loss = 0.0
    
    for data, labels in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()*data.size(0)
    # print training statistics 
    # calculate average loss over an epoch
    train_loss = train_loss/len(train_loader.dataset)

    print('Epoch: {} \tTraining Loss: {:.6f}'.format(
        epoch+1, 
        train_loss
        ))
    
# TEST THE TRAINED NETWORK
# initialize list to monitor test loss and accuracy    
test_loss = 0.0
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))

model.eval() # set the model to eval mode

for data, labels in test_loader:
    output = model(data)
    loss = criterion(output, labels)
    test_loss += loss.item()*data.size(0)
    # convert output probabilities to predicted class
    _, pred = torch.max(output, 1) # returns the max value and index along the rows.
    correct = np.squeeze(pred.eq(labels.data.view_as(pred)))
    # calculate the test accuracy
    for i in range(batch_size):
        label = labels.data[i]
        class_correct[label] += correct[i].item()
        class_total[label] += 1

# calculate and print avg test loss
test_loss = test_loss/len(test_loader.dataset)
print('Test Loss: {:.6f}\n'.format(test_loss))

for i in range(10):
    if class_total[i] > 0:
        print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (str(i), 100 * class_correct[i] / class_total[i],
                                            np.sum(class_correct[i]), np.sum(class_total[i])))
    else:
        print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))

print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (100. * np.sum(class_correct) / np.sum(class_total),
                        np.sum(class_correct), np.sum(class_total)))    

# Visualize Sample Test Results
# obtain one batch of test images
dataiter = iter(test_loader)
images, labels = dataiter.next()

# get sample outputs
output = model(images)
# convert output probabilities to predicted class
_, preds = torch.max(output, 1)
# prep images for display
images = images.numpy()

# plot the images in the batch, along with predicted and true labels
fig = plt.figure(figsize=(25, 4))
for idx in np.arange(20):
    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
    ax.imshow(np.squeeze(images[idx]), cmap='gray')
    ax.set_title("{} ({})".format(str(preds[idx].item()), str(labels[idx].item())),
                 color=("green" if preds[idx]==labels[idx] else "red"))

# Save the trained model
checkpoint = {'input_size': 784,
              'output_size': 10,
              'hidden_layers': [128, 64],
              'state_dict': model.state_dict()
              }

torch.save(checkpoint, 'mnist-nn.pth')

# Code to load the saved model
def load_checkpoint(filepath):
    checkpoint = torch.load(filepath)
    model = Net(checkpoint['input_size'], checkpoint['output_size'], checkpoint['hidden_layers'])
    model.load_state_dict(checkpoint['state_dict'])
    return model

loaded_model = load_checkpoint('mnist-nn.pth')

