import torch
import torchvision
from torchvision.datasets import MNIST

dataset = MNIST(root='data/', download = True)

len(dataset)

test_dataset = MNIST(root='data/', train = False)
len(test_dataset)

import matplotlib.pyplot as plt
%matplotlib inline

image, label = dataset[10]
plt.imshow(image)
print('Label: ',label)

import torchvision.transforms as transforms
dataset = MNIST(root='data/', train = True, transform = transforms.ToTensor())
img_tensor, label = dataset[0]
print(img_tensor.shape,label)

import numpy as np
def split_indices(n, val_pct):
    n_val = int(val_pct*n)
    idxs = np.random.permutation(n)
    print(idxs)
    return idxs[n_val:],idxs[:n_val]
train_indices, val_indices = split_indices(len(dataset), val_pct = 0.2)
print(len(train_indices), len(val_indices))

from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data.dataloader import DataLoader
import torch.nn.functional as F 
batch_size = 100

train_sampler = SubsetRandomSampler(train_indices)
train_loader = DataLoader(dataset, batch_size, sampler=train_sampler)

val_sampler = SubsetRandomSampler(val_indices)
val_loader = DataLoader(dataset,batch_size, sampler = val_sampler)

import torch.nn as nn 
input_size = 28*28
num_classes = 10
model = nn.Linear(input_size, num_classes)

loss_fn = F.cross_entropy
learning_rate = 0.001
optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate)

class MnistModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(input_size, num_classes)
        
    def forward(self, xb):
        xb = xb.reshape(-1, 784)
        out = self.linear(xb)
        return out
model = MnistModel()
def loss_batch(model, loss_func, xb, yb, opt=None, metric=None):
    # Calculate loss
    preds = model(xb)
    loss = loss_func(preds, yb)
                     
    if opt is not None:
        # Compute gradients
        loss.backward()
        # Update parameters             
        opt.step()
        # Reset gradients
        opt.zero_grad()
    
    metric_result = None

    if metric is not None:
        # Compute the metric
        metric_result = metric(preds, yb)
    
    return loss.item(), len(xb), metric_result

def evaluate(model, loss_fn, valid_dl, metric=None):
    with torch.no_grad():
        # Pass each batch through the model
        results = [loss_batch(model, loss_fn, xb, yb, metric=metric)
                   for xb,yb in valid_dl]
        # Separate losses, counts and metrics
        losses, nums, metrics = zip(*results)
        # Total size of the dataset
        total = np.sum(nums)
        # Avg. loss across batches 
        avg_loss = np.sum(np.multiply(losses, nums)) / total
        avg_metric = None
        if metric is not None:
            # Avg. of metric across batches
            avg_metric = np.sum(np.multiply(metrics, nums)) / total
    return avg_loss, total, avg_metric

def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.sum(preds == labels).item() / len(preds)

val_loss, total, val_acc = evaluate(model, loss_fn, val_loader, metric=accuracy)
print('Loss: {:.4f}, Accuracy: {:.4f}'.format(val_loss, val_acc))

def fit(epochs, model, loss_fn, opt, train_dl, valid_dl, metric=None):
    for epoch in range(epochs):
        # Training
        for xb,yb in train_dl:
            loss,_,_ = loss_batch(model, loss_fn, xb, yb, opt)

        # Evaluation
        result = evaluate(model, loss_fn, valid_dl, metric)
        val_loss, total, val_metric = result
        
        # Print progress
        if metric is None:
            print('Epoch [{}/{}], Loss: {:.4f}'
                  .format(epoch+1, epochs, val_loss))
        else:
            print('Epoch [{}/{}], Loss: {:.4f}, {}: {:.4f}'
                  .format(epoch+1, epochs, val_loss, metric.__name__, val_metric))

# Redifine model and optimizer

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

fit(5, model, F.cross_entropy, optimizer, train_loader, val_loader, accuracy)

# Replace these values with your results
accuracies = [0.1076, 0.6486, 0.7362, 0.7737, 0.7925, 0.8069, 
               0.8165, 0.8227, 0.8269, 0.8325, 0.8367,
               0.8399, 0.8438, 0.8463, 0.8482, 0.8512,
               0.8529, 0.8552, 0.8563, 0.8587, 0.8606]
plt.plot(accuracies, '-x')
plt.xlabel('epoch')
plt.ylabel('accuracy')
plt.title('Accuracy vs. No. of epochs');

torch.save(model.state_dict(),'mnist-logistic.pth')
model2 = MnistModel()
model2.load_state_dict(torch.load('mnist-logistic.pth'))
test_dataset = MNIST(root='data/', train = False, transform = transforms.ToTensor())
test_loader = DataLoader(test_dataset, batch_size = 100)
test_loss, total, test_acc = evaluate(model, loss_fn, test_loader, metric = accuracy)
print('LOSS:{:.4F}, Accuracy:{:.4f}'.format(test_loss, test_acc))
test_loss, total, test_acc = evaluate(model2, loss_fn, test_loader, metric = accuracy)
print('LOSS:{:.4F}, Accuracy:{:.4f}'.format(test_loss, test_acc))

import jovian
jovian.commit()