%load_ext autoreload
%autoreload 2
%matplotlib inline

import torch
from torch import tensor
from torch import nn
import torch.nn.functional as F
from torch.utils import data 
import matplotlib.pyplot as plt

from pathlib import Path
from IPython.core.debugger import set_trace
from fastai import datasets
from fastai.metrics import accuracy
import pickle, gzip, math, torch
import operator

MNIST_URL='http://deeplearning.net/data/mnist/mnist.pkl'

def get_data():
    path = datasets.download_data(MNIST_URL, ext='.gz')
    with gzip.open(path, 'rb') as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')
    return map(tensor, (x_train, y_train, x_valid, y_valid))

x_train, y_train, x_valid, y_valid = get_data()

class Dataset(data.Dataset):
    def __init__(self, x, y):
        self.x, self.y = x, y
    def __len__(self):
        return len(self.x)
    def __getitem__(self, i):
        return self.x[i], self.y[i]

# x = data, m = mean, s = standard deviation
def normalize(x, m, s): 
    return (x-m)/s

n, m = x_train.shape
c = (y_train.max()+1).numpy()
n, m, c

train_mean, train_std = x_train.mean(), x_train.std()
x_train = normalize(x_train, train_mean, train_std)
x_valid = normalize(x_valid, train_mean, train_std)

# batch size
bs = 64
train_ds, valid_ds = Dataset(x_train, y_train), Dataset(x_valid, y_valid)
train_dl, valid_dl = data.DataLoader(train_ds, bs, shuffle=True), data.DataLoader(valid_ds, bs, shuffle=False)

# learning rate
lr = 0.03
epoch = 10
nh = 50

def get_model():
    # loss function
    loss_func = F.cross_entropy
    model = nn.Sequential(nn.Linear(m, nh), nn.ReLU(), nn.Linear(nh,c))
    return model, loss_func

def fit():
    losses, metrics = [], []
    # e = epoch number
    for e in range(epoch):
        for xb, yb in train_dl:

            # Feedforward
            yhatb = model(xb)
            loss = loss_func(yhatb, yb)

            # Metrics
            acc = accuracy(yhatb, yb)
            losses.append(loss); metrics.append(acc)

            # Backpropagation
            loss.backward()

            optim.step()
            optim.zero_grad()

    plot_metrics(losses, metrics)

def plot_metrics(losses, metrics):
    x = torch.arange(len(losses)).numpy()
    fig,ax = plt.subplots(figsize=(9, 9))
    ax.grid(True)
    ax.plot(x, losses, label="Loss")
    ax.plot(x, metrics, label="Accuracy")
    ax.legend(loc='upper right')

def fit2():
    # e = epoch number
    for e in range(epoch):

        # Set Model in Train Mode
        model.train()

        for xb, yb in train_dl:
            yhatb = model(xb)
            loss = loss_func(yhatb, yb)
            loss.backward()
            optim.step()
            optim.zero_grad()

        # Set Model in Evaluation Mode
        model.eval()

        # Metrics
        with torch.no_grad():
            # tot_loss = total loss, tot_acc = total accuracy
            tot_loss, tot_acc = 0., 0.
            for xb, yb in valid_dl:
                yhatb = model(xb)
                tot_acc += accuracy(yhatb, yb)
                tot_loss += loss_func(yhatb, yb)
            # nv = number of validation batch
            nv = len(valid_ds)/bs
            print(f'epoch={e}, valid_loss={tot_loss/nv}, valid_acc={tot_acc/nv}')            
    return tot_loss/nv, tot_acc/nv
    

model, loss_func = get_model()
optim = torch.optim.SGD(model.parameters(), lr=lr)
fit2()