%load_ext autoreload
%autoreload 2
%matplotlib inline

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

from pathlib import Path
from IPython.core.debugger import set_trace
from fastai import datasets
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()

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

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)

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

class Model(nn.Module):
    # number of input feature, number of hidden feature, number of output feature
    def __init__(self, n_in, nh, n_out):
        super().__init__()
#         set_trace()
        self.layers = [nn.Linear(n_in, nh), nn.ReLU(), nn.Linear(nh, n_out)]        
    def __call__(self, x):
        for l in self.layers:
            x = l(x)
        return x

def get_model():
    # loss function
    loss_func = F.cross_entropy
    model = Model(m, nh, c.numpy())
    return model, loss_func

def accuracy(yhat, y):
    return (torch.argmax(yhat, dim=-1) == y).float().mean()

# batch size
bs = 64
# learning rate
lr = 0.03
epoch = 1
nh = 50

model, loss_func = get_model()
losses, metrics = [], []

def fit():
    losses, metrics = [], []
    # e = epoch number
    for e in range(epoch):
        # b = batch number
        for b in range((n-1)//bs + 1):

            # data
            start_b = b*bs
            end_b = start_b+bs
            xb, yb = x_train[start_b:end_b], y_train[start_b:end_b]

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

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

            # Backpropagation
            loss.backward()

            # Update Wight and Bias with SGD
            with torch.no_grad():
                for l in model.layers:
                    if hasattr(l, 'weight'):
                        l.weight -= l.weight.grad * lr
                        l.bias   -= l.bias.grad * lr 
                        l.weight.grad.zero_()
                        l.bias.grad.zero_()       
                        
    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')

model, loss_func = get_model()
fit()

class 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]

train_ds, valid_ds = Dataset(x_train, y_train), Dataset(x_valid, y_valid)

xb, yb = train_ds[0:100]
xb.shape, yb.shape

plt.imshow(xb[0].view(28,28))

def fit2():
    losses, metrics = [], []
    # e = epoch number
    for e in range(epoch):
        # b = batch number
        for b in range((n-1)//bs + 1):

            # data
            start_b = b*bs
            end_b = start_b+bs            
            xb, yb = train_ds[start_b:end_b]

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

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

            # Backpropagation
            loss.backward()

            # Update Wight and Bias with SGD
            with torch.no_grad():
                for l in model.layers:
                    if hasattr(l, 'weight'):
                        l.weight -= l.weight.grad * lr
                        l.bias   -= l.bias.grad * lr 
                        l.weight.grad.zero_()
                        l.bias.grad.zero_()      
                        
    plot_metrics(losses, metrics)

model, loss_func = get_model()
fit2()

class DataLoader():
    # ds = Dataset, bs = batch size
    def __init__(self, ds, bs):
        self.ds, self.bs = ds, bs
    def __iter__(self):
        for i in range(0, len(self.ds), self.bs):
            yield self.ds[i:i+self.bs]

train_dl, valid_dl = DataLoader(train_ds, bs), DataLoader(valid_ds, bs)

train_iter = iter(train_dl)

xb, yb = next(train_iter)
plt.imshow(xb[0].view(28,28))

xb, yb = next(train_iter)
plt.imshow(xb[0].view(28,28))

def fit3():
    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()

            # Update Wight and Bias with SGD
            with torch.no_grad():
                for l in model.layers:
                    if hasattr(l, 'weight'):
                        l.weight -= l.weight.grad * lr
                        l.bias   -= l.bias.grad * lr 
                        l.weight.grad.zero_()
                        l.bias.grad.zero_()       

    plot_metrics(losses, metrics)

model, loss_func = get_model()
fit3()