"""
# Proyecto 1 Computación Emergente
Integrantes:
- Samuel Mariña
- María G. Cafarelli
"""
"""
## Carga de paquetes
"""
import numpy as np
import pandas as pd

import torch
from torch.utils.data import Dataset,DataLoader
from torchvision import transforms
from torch import optim
import torch.nn as nn
import torch.nn.functional as F

from pathlib import Path
from sklearn.model_selection import train_test_split

import matplotlib.pyplot as plt
%matplotlib inline

import time
from time import process_time
"""
## Definición de Variables
"""
batch_size = 300
epochs = 30
test_size = 0.15
learning_rate = 0.0001
"""
## Carga de Data (Training Set)
"""
# Imágenes y labels de entrenamiento (Training Set)
train_images = np.load("../input/kuzushiji/kmnist-train-imgs.npz")["arr_0"]
train_labels = np.load("../input/kuzushiji/kmnist-train-labels.npz")["arr_0"]
"""
## Separación de Datos
"""
print("Tamaño Set de Entrenamiento antes:" ,train_images.shape)

train_images, test_images, train_labels, test_labels = train_test_split(train_images, train_labels, test_size = test_size)
# #Imprimimos para confirmar
print("Tamaño Set de Entrenamiento después: ", train_images.shape)
print("Tamaño Set de Validacion: ", test_images.shape)
"""
## Normalización de Datos
"""
def normalizar(dataset, labels):
    data = dataset.astype('float32')
    data /= 255
    data = np.reshape(data, (len(data), 1, 28, 28))
    newLabels = labels.astype('int64')
    return data, newLabels
train_images, train_labels = normalizar(train_images, train_labels)
test_images, test_labels = normalizar(test_images, test_labels)
"""
## Conversión a Tensores
"""
x_train_images = torch.Tensor(train_images)
y_train_images = torch.Tensor(train_labels).type(torch.LongTensor)
x_valid_images = torch.Tensor(test_images)
y_valid_images = torch.Tensor(test_labels).type(torch.LongTensor)

print("Set de Entrenamiento: " + str(x_train_images.shape))
print("Labels del Set de Entrenamiento: " + str(y_train_images.shape))
print("Set de Validación: " + str(x_valid_images.shape))
print("Labels del Set de Validación: " + str(y_valid_images.shape))
"""
## Agregar Device si se tiene GPU
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
"""
## Obtener Datasets
"""
# Training Dataset
train_dataset = torch.utils.data.TensorDataset(x_train_images, y_train_images)
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False)

#Validation Dataset
valid_dataset = torch.utils.data.TensorDataset(x_valid_images, y_valid_images)
validloader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
"""
# Arquitectura del Modelo
"""
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1), # output dim = (28x28x32)
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Dropout(p=0.2)
        )
        
        self.layer2 = nn.Sequential(
            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1), # output dim = (28x28x32)
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Dropout(p=0.2)
        )
        
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=0), # output = (26x26x64)
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2), # output = (13x13x64)
            nn.Dropout(p=0.2)
        )
        
        
        self.fc1 = nn.Sequential(
            nn.Linear(10816, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
            nn.LogSoftmax(dim=1)
        )
        
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        
        return x
"""
# Entrenamiento de la Red
"""
model = CNN()
model.cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

training_loss = []
valid_loss = []
training_accuracy = []
valid_accuracy = []
trainingTimeStart = time.process_time()
for epoch in range(epochs):
    model.train()
    running_loss = 0
    accuracy = 0
    steps = 0
    
    for images, labels in trainloader:    
        images = images.cuda()
        labels = labels.cuda()
        #Training pass
        output = model(images)
        loss = criterion(output, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        
        _, predicted = torch.max(output, 1)
        accuracy += (predicted == labels).sum()
        steps += 1
        
    training_loss.append(running_loss/steps)
    training_accuracy.append(100 * accuracy.cpu().numpy()/len(train_dataset))
    model.eval()
    
    with torch.no_grad():
        iter_loss = 0
        accuracy = 0
        steps = 0
        
        for images, labels in validloader:
            images = images.cuda()
            labels = labels.cuda()
            output = model(images)
            loss = criterion(output, labels)
            
            iter_loss += loss.item()
            _, predicted = torch.max(output, 1)
            accuracy += (predicted == labels).sum()
            steps += 1
            
        valid_loss.append(iter_loss/steps)
        valid_accuracy.append(100 * accuracy.cpu().numpy()/len(valid_dataset))
        
    print ('Epoch {}/{}, Training Loss: {:.3f}, Training Accuracy: {:.2f}%, Validation Loss: {:.3f}, Validation Acc: {:.2f}%'
            .format(epoch+1, epochs, training_loss[-1], training_accuracy[-1], valid_loss[-1], valid_accuracy[-1]))
    
trainingTimeStop = time.process_time() - trainingTimeStart
print('Tiempo de ejecución: ' + str(trainingTimeStop))
plt.plot(training_loss, color='red', label="Training loss")
plt.plot(valid_loss, color='blue',label="Validation loss")
plt.legend()
plt.show()
plt.plot(training_accuracy, color='red', label="Training acc")
plt.plot(valid_accuracy, color='blue',label="Validation acc")
plt.legend()
plt.show()
"""
# Resultado
"""
testing_images = np.load("../input/kuzushiji/kmnist-test-imgs.npz")["arr_0"]
testing_labels = np.load("../input/kuzushiji/kmnist-test-labels.npz")["arr_0"]
new_testing_images, new_testing_labels = normalizar(testing_images, testing_labels)
x_testing_images = torch.Tensor(new_testing_images)
y_testing_images = torch.Tensor(new_testing_labels).type(torch.LongTensor)
print("Set de Prueba: " + str(x_testing_images.shape))
print("Labels del Set de Prueba: " + str(y_testing_images.shape))
test_dataset = torch.utils.data.TensorDataset(x_testing_images, y_testing_images)
testloader = torch.utils.data.DataLoader(test_dataset, batch_size=300, shuffle=False)
model.eval()
testingTimeStart = time.process_time()
with torch.no_grad():
    accuracy = 0
    total = 0
    test_loss = 0
    
    for images, labels in testloader:
        images = images.cuda()
        labels = labels.cuda()
        output = model(images)
        loss = criterion(output, labels)
        test_loss += loss.item()
        test_loss = test_loss/len(testloader)
        
        _, predicted = torch.max(output.data, 1)
        total += labels.size(0)
        accuracy += (predicted == labels).sum().item()
        test_accuracy = (100 * accuracy)/total
    
    print('Test Loss: {:.4f}, Test Accuracy: {:.2f}%'.format(test_loss, test_accuracy))
testingTimeStop = time.process_time() - testingTimeStart
print("Tiempo de ejecución: " + str(testingTimeStop))