import numpy as np
from sklearn.datasets import load_iris

iris = load_iris()
data, labels = iris.data[:,0:3], iris.data[:,3]

"""
Explanation: Introducing PyTorch
written by Gene Kogan, updated by Sebastian Quinard

In the next cell, we introduce PyTorch, which is an open-source framework which impelments machine learning methodology, particularly that of deep neural networks, by optimizing the efficiency of the computation. We do not have to deal so much with the details of this. Most importantly, PyTorch efficiently implement backpropagation to train neural networks on the GPU.
To start, we will re-implement what we did in the last notebook, a neural network to to predict the sepal width of the flowers in the Iris dataset. In the last notebook, we trained a manually-coded neural network for this, but this time, we'll use PyTorch instead. 
Let's load the Iris dataset again.
End of explanation
"""


num_samples = len(labels)  # size of our dataset
shuffle_order = np.random.permutation(num_samples)
data = data[shuffle_order, :]
labels = labels[shuffle_order]

# normalize data and labels to between 0 and 1 and make sure it's float32
data = data / np.amax(data, axis=0)
data = data.astype('float32')
labels = labels / np.amax(labels, axis=0)
labels = labels.astype('float32')

# print out the data
print("shape of X", data.shape)
print("first 5 rows of X\n", data[0:5, :])
print("first 5 labels\n", labels[0:5])

"""
Explanation: First we need to shuffle and pre-process the data. Pre-processing in this case is normalization of the data, as well as converting it to a properly-shaped numpy array.
End of explanation
"""


# let's rename the data and labels to X, y
X, y = data, labels

test_split = 0.3  # percent split

n_test = int(test_split * num_samples)

x_train, x_test = X[n_test:, :], X[:n_test, :]
x_train = torch.from_numpy(x_train)
x_test = torch.from_numpy(x_test)
y_train, y_test = y[n_test:], y[:n_test] 
y_train = torch.from_numpy(y_train)
y_test = torch.from_numpy(y_test)
print('%d training samples, %d test samples' % (x_train.shape[0], x_test.shape[0]))

"""
Explanation: Overfitting and validation
In our previous guides, we always evaluated the performance of the network on the same data that we trained it on. But this is wrong; our network could learn to "cheat" by overfitting to the training data (like memorizing it) so as to get a high score, but then not generalize well to actually unknown examples.
In machine learning, this is called "overfitting" and there are several things we have to do to avoid it. The first thing is we must split our dataset into a "training set" which we train on with gradient descent, and a "test set" which is hidden from the training process that we can do a final evaluation on to get the true accuracy, that of the network trying to predict unknown samples.
Let's split the data into a training set and a test set. We'll keep the first 30% of the dataset to use as a test set, and use the rest for training.
End of explanation
"""


import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(3, 8) # we get 3 from input dimension, and 8 from desired output
        self.fc2 = nn.Linear(8, 1) # Output Layer
    
    def forward(self, x):
        x = F.sigmoid(self.fc1(x))
        x = self.fc2(x)
        return x

"""
Explanation: Creating the Model
In PyTorch, to instantiate a neural network model, we inherit the class nn.Module which grants our class Net all the functionality of the nn.Module. We then instantiate it with the init() class. Note however, we must also instantiate the class we are inheriting from (i.e. nn.Module). This alone creates an empty neural network, so to populate it, we add the type of layer we want. In this case, we add a linear layer, a layer that is \"fully-connected,\" meaning all of its neurons are connected to all the neurons in the previous layer, with no empty connections. This may seem confusing at first because we have not yet seen neural network layers which are not fully-connected; we will see this in the next chapter when we introduce convolutional networks.
Next we see the addition of the forward method. We couple the forward method with the layer(s) above to apply a variety of activation functions onto the specified layer.
Finally, we will add the output layer, which will be a fully-connected layer whose size is 1 neuron. This neuron will contain our final output.
End of explanation
"""


net = Net()

"""
Explanation: That may be a lot to take in, but once you fully understand the excerpt above, this structure will be used time and time again to build increasingly complex neural networks.
Next we instantiate a new object based on the class.
End of explanation
"""


print(net)

"""
Explanation: We can also get a readout of the current state of the network using print(net):
End of explanation
"""


from torch import optim
optimizer = optim.SGD(net.parameters(), lr=0.01)
criterion = nn.MSELoss()

"""
Explanation: So we've added 9 parameters, 8x1 weights between the hidden and output layers, and 1 bias in the output. So we have 41 parameters in total.
Now we are finished specifying the architecture of the model. Now we need to specify our loss function and optimizer, and then compile the model. Let's discuss each of these things.
First, we specify the loss. The standard for regression, as we said before is sum-squared error (SSE) or mean-squared error (MSE). SSE and MSE are basically the same, since the only difference between them is a scaling factor ($\frac{1}{n}$) which doesn't depend on the final weights.
The optimizer is the flavor of gradient descent we want. The most basic optimizer is "stochastic gradient descent" or SGD which is the learning algorithm we have used so far. We have mostly used batch gradient descent so far, which means we compute our gradient over the entire dataset. For reasons which will be more clear when we cover learning algorithms in more detail, this is not usually favored, and we instead calculate the gradient over random subsets of the training data, called mini-batches.
Once we've specified our loss function and optimizer, the model is compiled.
End of explanation
"""



def fullPass(data, labels):

    running_loss = 0.0
    for i in range(0,data.size()[0]):
        optimizer.zero_grad()
        outputs = net(data[i])
        loss = criterion(outputs, labels[i])
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss
        if i % data.size()[0] == data.size()[0]-1:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / data.size()[0]))
            running_loss = 0.0

net.train()
for epoch in range(400):
    fullPass(x_train, y_train);


print('Finished Training')


"""
Explanation: We are finally ready to train. First we must zero our gradients, so as not to rely on the previously uncovered gradient in our solution. The opposite is imperitive to RNN, as we need the last result to influence the next.
Next we create a forward pass on our neural network. The loss function is then applied to determine the level of error. After, we complete a backwards pass to compute the new weights.  
Loss is also printed.
End of explanation
"""


net.eval()
fullPass(x_test, y_test)

y_pred = net(x_test)

"""
Explanation: As you can see above, we train our network down to a validation MSE < 0.01. Notice that both the training loss ("loss") and validation loss ("val_loss") are reported. It's normal for the training loss to be lower than the validation loss, since the network's objective is to predict the training data well. But if the training loss is much lower than our validation loss, it means we are overfitting and may not expect to receive very good results.
We can evaluate the training set one last time at the end using eval.
End of explanation
"""


def MSE(y_pred, y_test):
    return (1.0/len(y_test)) * np.sum([((y1[0]-y2)**2) for y1, y2 in list(zip(y_pred, y_test))])

print("MSE is %0.4f" % MSE(y_pred, y_test))

"""
Explanation: We can manually calculate MSE as a sanity check:
End of explanation
"""


x_sample = x_test[0].reshape(1, 3)   # shape must be (num_samples, 3), even if num_samples = 1
y_prob = net(x_sample)

print("predicted %0.3f, actual %0.3f" % (y_prob[0][0], y_test[0]))

"""
Explanation: We can also predict the value of a single unknown example or a set of them in th following way:
End of explanation
"""


