# -*- coding: utf-8 -*-
"""hw11_pytorch_21600004_KangSeokUn.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1BQnNOTFy9jlg-NA5FYrvl1ThZKnFW09K
"""

from __future__ import print_function
import torch
import torchvision            
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets
from torchvision import transforms

import time
import math
import random as rd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import datetime
from matplotlib.pyplot import imshow, imsave
print(torch.__version__)

transform = transforms.Compose(
    [transforms.ToTensor(),                               # image to tensor
     transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # normalize to "(x-mean)/std"
    ])

# Download the MNIST Dataset
mnist_train = datasets.MNIST(root='../data/', train=True, transform=transform, download=True)
mnist_test = datasets.MNIST(root='../data/', train=False, transform=transform, download=True)

batch_size = 64

train_loader = DataLoader(dataset=mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)
test_loader = DataLoader(dataset=mnist_test, batch_size=100, shuffle=False, drop_last=False)

def timeSince(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)

MODEL_NAME = 'MLP'
# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DEVICE = "cpu"
print("MODEL_NAME = {}, DEVICE = {}".format(MODEL_NAME, DEVICE))

class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.fc1 = nn.Linear(784, 200)
    self.fc2 = nn.Linear(200, 10)
    
    
  def forward(self, x):
    x = x.view(x.size(0), -1)
    x = torch.sigmoid(self.fc1(x))
    x = self.fc2(x)

    return F.log_softmax(x, dim=1)

# set loss function and optimizer
model = Net().to(DEVICE)

criterion = F.nll_loss
optim = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

max_epoch = 4        # maximum number of epochs
step = 0             # initialize step counter variable

plot_every = 100
total_loss = 0 # Reset every plot_every iters
all_losses = [] # reset loss history

start = time.time()

for epoch in range(max_epoch):
    for idx, (images, labels) in enumerate(train_loader):
        # Training Discriminator
        x, y = images.to(DEVICE), labels.to(DEVICE) # (N, 1, 28, 28), (N, )
        
        y_hat = model(x) # (N, 10)  # forward propagation
       
        loss = criterion(y_hat, y)  # computing loss
        loss_val = loss.item()
          
        optim.zero_grad()           # reset gradient
        loss.backward()             # back-propagation (compute gradient)
        optim.step()                # update parameters with gradient
        
        # periodically print loss
        if step % 10 == 0:
            print('Epoch({}): {}/{}, Step: {}, Loss: {}'.format(timeSince(start), epoch, max_epoch, step, loss.item()))
        
        
        all_losses.append(loss_val)
        
        # periodically evalute model on test data
        if step % 20 == 0:
            model.eval()
            acc = 0.
            with torch.no_grad():   # disable autograd
                for idx, (images, labels) in enumerate(test_loader):
                    x, y = images.to(DEVICE), labels.to(DEVICE) # (N, 1, 28, 28), (N, )
                    y_hat = model(x) # (N, 10)
                    loss = criterion(y_hat, y)
                    _, indices = torch.max(y_hat, dim=-1)     # find maxmum along the last axis (argmax of each row)
                                                              # ex) max_value, max_idx = torch.max(input, dim)
                    acc += torch.sum(indices == y).item()     # count correctly classified samples
                                                              # torch.sum() returns Tensor. Tensor.item() converts it to a value
            print('*'*20, 'Test', '*'*20)
            print('Step: {}, Loss: {}, Accuracy: {} %'.format(step, loss.item(), acc/len(mnist_test)*100))
            print('*'*46)
            model.train()           # turn to train mode (enable autograd)
        step += 1
        if(step >= 3000):
            break
    if(step >= 3000):
        break
        
print("---train {}s seconds---".format(time.time()-start))

plt.figure()
plt.plot(all_losses, c='r', label="Cost")
plt.legend()
plt.title('Cost per epoch')
plt.savefig("./pytorch.png")
plt.clf()

# Test
model.eval()
acc = 0.
with torch.no_grad():
    for idx, (images, labels) in enumerate(test_loader):
        x, y = images.to(DEVICE), labels.to(DEVICE) # (N, 1, 28, 28), (N, )
        y_hat = model(x) # (N, 10)
        loss = criterion(y_hat, y)
        _, indices = torch.max(y_hat, dim=-1)
        acc += torch.sum(indices == y).item()
print('*'*20, 'Test', '*'*20)
print('Step: {}, Loss: {}, Accuracy: {} %'.format(step, loss.item(), acc/len(mnist_test)*100))
print('*'*46)

# select the random
print(len(mnist_test))
idx_set = [ rd.randint(0, 10000 + 1) for idx in range(10)]

for idx in idx_set:
  img, y = mnist_test[idx]


  sample = img.unsqueeze(dim=0).to(DEVICE)
  out = model(sample).cpu()
  _, idx = out.max(dim=-1)
  print("prediction = {} | Label = {}".format(idx.item(), y))
  # imshow(img[0], cmap='gray')
  plt.imshow(img[0], cmap='gray')
  plt.show()

