import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from torchvision import datasets
import torch
import torch.nn as nn
import torch.optim as optim

"""
Explanation: FashionMNIST classification with Multilayer perceptrons
https://github.com/zalandoresearch/fashion-mnist
Exercises

Try changing the model and training parameters (hidden size, epoch, batch size).
Add another hidden layer to the model.
All input values are between 0 to 255 (RGB brightness values). Normalize them (scale them to between 0 and 1). You should see a small increase in accuracy.
Try standardizing the input values (mean 0, std 1).
Implement early stopping. Instead of a fix number of epochs, train until the model stops improving on the dev set. If the dev loss stop decreasing for a few epochs, training should stop.
We plot one random sample and its label. Plot a random sample from each class on 10 subplots.
Limit printing the loss and the accuracy to every 10th epoch instead of every epoch. 
ADVANCED: replace the network with a convolutional neural network. Convolutions should work on 2D images, change the preprocessing steps accordingly. You should see a big increase in accuracy.
End of explanation
"""


train_data = datasets.FashionMNIST('data', download=True, train=True)
# we need FloatTensors as input
train_X = train_data.data.float()
train_y = train_data.targets
test_data = datasets.FashionMNIST('data', download=True, train=False)
test_X = test_data.data.float()
test_y = test_data.targets

"""
Explanation: Preparing the dataset
Download and convert to float
End of explanation
"""


labels = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

idx = np.random.randint(len(train_X))
sample_X = train_X[idx].numpy()
sample_y = train_y[idx].numpy()
print("Label: {}".format(labels[sample_y]))
plt.imshow(sample_X, 'gray')

"""
Explanation: IMPORTANT: Inspect the dataset manually
We will plot a random sample with its labels:
End of explanation
"""


train_X.size()

"""
Explanation: Understanding the sizes of the data
Make sure you always understand the output of tensor.size()
Train size X image_x X image_y
End of explanation
"""


np.unique(train_y.numpy(), return_counts=True)

"""
Explanation: Are the classes equally distributed?
End of explanation
"""


print("Before flattening:")
print("Train size:", train_X.size(), train_y.size())
print("Test size:", test_X.size(), test_y.size())

train_X = train_X.view(-1, 28 * 28).squeeze(1)
test_X = test_X.view(-1, 28 * 28).squeeze(1)

print("\nAfter flattening:")
print("Train size:", train_X.size(), train_y.size())
print("Test size:", test_X.size(), test_y.size())

"""
Explanation: Flatten images
End of explanation
"""


all_idx = np.arange(len(train_X))
np.random.shuffle(all_idx)
train_idx = all_idx[:50000]
dev_idx = all_idx[50000:]
print("The overlap between train and dev should be an empty set:", set(train_idx) & set(dev_idx))
print("")

dev_X = train_X[dev_idx]
dev_y = train_y[dev_idx]
train_X = train_X[train_idx]
train_y = train_y[train_idx]
print("Train size:", train_X.size(), train_y.size())
print("Dev size:", dev_X.size(), dev_y.size())
print("Test size:", test_X.size(), test_y.size())

"""
Explanation: Splitting the train set into train and dev set
End of explanation
"""


class BatchedIterator:
    def __init__(self, X, y, batch_size):
        self.X = X
        self.y = y
        self.batch_size = batch_size
    
    def iterate_once(self):
        for start in range(0, len(self.X), self.batch_size):
            end = start + self.batch_size
            yield self.X[start:end], self.y[start:end]

"""
Explanation: Helper class that creates minibatches
End of explanation
"""


train_iter = BatchedIterator(train_X, train_y, 33333)

for batch in train_iter.iterate_once():
    print(batch[0].size(), batch[1].size())

"""
Explanation: Testing the iterator:
End of explanation
"""


class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim):
        super().__init__()
        self.input_layer = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.output_layer = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, X):
        h = self.input_layer(X)
        h = self.relu(h)
        out = self.output_layer(h)
        return out

"""
Explanation: Neural network
Definition
End of explanation
"""


model = SimpleClassifier(
    input_dim=train_X.size(1),
    output_dim=10,
    hidden_dim=50
)
model

for n, p in model.named_parameters():
    print(n, p.size())

"""
Explanation: Model instance
End of explanation
"""


criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

"""
Explanation: Loss function and optimizer
End of explanation
"""


test_pred = model(test_X).max(axis=1)[1]
test_acc = torch.eq(test_pred, test_y).sum().float() / len(test_X)
test_acc

"""
Explanation: Sanity check
The model should perform close to random at this point.
Note we compute accuracy manually. Make sure you understand it.
End of explanation
"""


batch_size = 1000
train_iter = BatchedIterator(train_X, train_y, batch_size)
dev_iter = BatchedIterator(dev_X, dev_y, batch_size)
test_iter = BatchedIterator(test_X, test_y, batch_size)

all_train_loss = []
all_dev_loss = []
all_train_acc = []
all_dev_acc = []

n_epochs = 10
for epoch in range(n_epochs):
    # training loop
    for bi, (batch_x, batch_y) in enumerate(train_iter.iterate_once()):
        
        y_out = model(batch_x)
        loss = criterion(y_out, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    # one train epoch finished, evaluate on the train and the dev set (NOT the test)
    train_out = model(train_X)
    train_loss = criterion(train_out, train_y)
    all_train_loss.append(train_loss.item())
    train_pred = train_out.max(axis=1)[1]
    train_acc = torch.eq(train_pred, train_y).sum().float() / len(train_X)
    all_train_acc.append(train_acc)
    
    dev_out = model(dev_X)
    dev_loss = criterion(dev_out, dev_y)
    all_dev_loss.append(dev_loss.item())
    dev_pred = dev_out.max(axis=1)[1]
    dev_acc = torch.eq(dev_pred, dev_y).sum().float() / len(dev_X)
    all_dev_acc.append(dev_acc)
    
    print(f"Epoch: {epoch}\n  train accuracy: {train_acc}  train loss: {train_loss}")
    print(f"  dev accuracy: {dev_acc}  dev loss: {dev_loss}")

"""
Explanation: Training loop
End of explanation
"""


test_pred = model(test_X).max(axis=1)[1]
test_acc = torch.eq(test_pred, test_y).sum().float() / len(test_X)
test_acc

"""
Explanation: Test accuarcy
End of explanation
"""


plt.plot(all_train_loss, label='train')
plt.plot(all_dev_loss, label='dev')
plt.legend()

plt.plot(all_train_acc, label='train')
plt.plot(all_dev_acc, label='dev')
plt.legend()

"""
Explanation: Plot loss functions, accuracies
End of explanation
"""


toy_X = train_X[:5]
toy_y = train_y[:5]

model = SimpleClassifier(
    input_dim=train_X.size(1),
    output_dim=10,
    hidden_dim=500
)
optimizer = optim.Adam(model.parameters())

batch_size = 20
toy_train_iter = BatchedIterator(toy_X, toy_y, batch_size)

all_train_loss = []
all_dev_loss = []
all_train_acc = []
all_dev_acc = []

n_epochs = 20
for epoch in range(n_epochs):
    # training loop
    for bi, (batch_x, batch_y) in enumerate(toy_train_iter.iterate_once()):
        
        y_out = model(batch_x)
        loss = criterion(y_out, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    # one train epoch finished, evaluate on the train and the dev set (NOT the test)
    train_out = model(toy_X)
    train_loss = criterion(train_out, toy_y)
    all_train_loss.append(train_loss.item())
    train_pred = train_out.max(axis=1)[1]
    train_acc = torch.eq(train_pred, toy_y).sum().float() / len(toy_X)
    all_train_acc.append(train_acc)
    
    dev_out = model(dev_X)
    dev_loss = criterion(dev_out, dev_y)
    all_dev_loss.append(dev_loss.item())
    dev_pred = dev_out.max(axis=1)[1]
    dev_acc = torch.eq(dev_pred, dev_y).sum().float() / len(dev_X)
    all_dev_acc.append(dev_acc)

fig, ax = plt.subplots(1, 2, figsize=(12, 5))
ax[0].set_title("Loss")
ax[1].set_title("Accuracy")
ax[0].set_xlabel("epoch")
ax[1].set_xlabel("epoch")
ax[0].plot(all_train_loss, label='train')
ax[0].plot(all_dev_loss, label='dev')
ax[1].plot(all_train_acc, label='train')
ax[1].plot(all_dev_acc, label='dev')
plt.legend()

"""
Explanation: Overfitting
We'll train a larger network on a much smaller set:
End of explanation
"""


class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, output_dim,
                 hidden_dim):
        super().__init__()
        self.input_layer = nn.Linear(
            input_dim, hidden_dim)
        
        # let's add some extra layers in a list
        self.extra_layers = [
            nn.Linear(hidden_dim, 100),
            nn.ReLU(),
            nn.Linear(100, hidden_dim),
            nn.ReLU(),
        ]
        
        self.relu = nn.ReLU()
        self.output_layer = nn.Linear(
            hidden_dim, output_dim)
        
    def forward(self, X):
        h = self.input_layer(X)
        h = self.relu(h)
        # passing through extra layers
        for layer in self.extra_layers:
            h = layer(h)
        out = self.output_layer(h)
        return out

"""
Explanation: Gotchas
List elements are not trained
Modules such as nn.Linear are not registered as parameters if they're not direct attributes of a module.
In this example we add two extra layers as a list to the module:
End of explanation
"""


m = SimpleClassifier(4, 5, 6)
print(m)

print("Parameters:")
for name, param in m.named_parameters():
    print("Name: {}, size: {}".format(name, param.size()))

"""
Explanation: they are not part of the model
End of explanation
"""


class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, output_dim,
                 hidden_dim):
        super().__init__()
        self.input_layer = nn.Linear(
            input_dim, hidden_dim)
        
        # use ModuleList
        self.extra_layers = nn.ModuleList([
            nn.Linear(hidden_dim, 100),
            nn.ReLU(),
            nn.Linear(100, hidden_dim),
            nn.ReLU(),
        ])
        
        self.relu = nn.ReLU()
        self.output_layer = nn.Linear(
            hidden_dim, output_dim)
        
    def forward(self, X):
        h = self.input_layer(X)
        h = self.relu(h)
        # passing through extra layers
        for layer in self.extra_layers:
            h = layer(h)
        out = self.output_layer(h)
        return out

m = SimpleClassifier(4, 5, 6)
print(m)

print("Parameters:")
for name, param in m.named_parameters():
    print("Name: {}, size: {}".format(name, param.size()))

"""
Explanation: One solution is nn.ModuleList
End of explanation
"""


use_cuda = torch.cuda.is_available()
print(use_cuda)

"""
Explanation: Using the GPU
End of explanation
"""


if use_cuda:
    model = model.cuda()
    criterion = criterion.cuda()

"""
Explanation: Moving things manually to the GPU:

model: move once
criterion: move once
data: move one batch at a time

This should be automatically handled by your code the following way:
End of explanation
"""


