import torch 
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as dsets
import matplotlib.pylab as plt
import numpy as np


"""
Explanation: <div class="alert alert-block alert-info" style="margin-top: 20px">
 <a href="http://cocl.us/pytorch_link_top"><img src = "http://cocl.us/Pytorch_top" width = 950, align = "center"></a>

<img src = "https://ibm.box.com/shared/static/ugcqz6ohbvff804xp84y4kqnvvk3bq1g.png" width = 200, align = "center">


<h1 align=center><font size = 5>Convolutional Neral Network </font></h1> 


# Table of Contents
In this lab, we will use a Convolutional Neral Networks to classify handwritten digits from the MNIST database

<div class="alert alert-block alert-info" style="margin-top: 20px">
<li><a href="#ref0">Helper functions </a></li>

<li><a href="#ref1"> Prepare Data </a></li>
<li><a href="#ref2">Convolutional Neral Network </a></li>
<li><a href="#ref3">Define Softmax , Criterion function, Optimizer and Train the  Model</a></li>
<li><a href="#ref4">Analyse Results</a></li>

<br>
<p></p>
Estimated Time Needed: <strong>25 min</strong>
</div>

<hr>

<a id="ref0"></a>
<h2 align=center>Helper functions </h2>
End of explanation
"""


def plot_channels(W):
    #number of output channels 
    n_out=W.shape[0]
    #number of input channels 
    n_in=W.shape[1]
    w_min=W.min().item()
    w_max=W.max().item()
    fig, axes = plt.subplots(n_out,n_in)
    fig.subplots_adjust(hspace = 0.1)
    out_index=0
    in_index=0
    #plot outputs as rows inputs as columns 
    for ax in axes.flat:
    
        if in_index>n_in-1:
            out_index=out_index+1
            in_index=0
              
        ax.imshow(W[out_index,in_index,:,:], vmin=w_min, vmax=w_max, cmap='seismic')
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        in_index=in_index+1

    plt.show()

"""
Explanation: <code>plot_channels</code>: 
plot out the kernel parameters of  each channel
End of explanation
"""


def plot_parameters (W,number_rows =1,name="",i=0):
    W=W.data[:,i,:,:]
    n_filters=W.shape[0]

    w_min=W.min().item()
    w_max=W.max().item()
    fig, axes = plt.subplots(number_rows, n_filters//number_rows)
    fig.subplots_adjust(hspace = 0.4)

    for i,ax in enumerate(axes.flat):
        if i< n_filters:
            # Set the label for the sub-plot.
            ax.set_xlabel( "kernel:{0}".format(i+1))

            # Plot the image.
            ax.imshow(W[i,:], vmin=w_min, vmax=w_max, cmap='seismic')
            ax.set_xticks([])
            ax.set_yticks([])
    plt.suptitle(name, fontsize=10)    
    plt.show()

"""
Explanation: <code>plot_parameters</code>: plot out the kernel parameters of  each channel
End of explanation
"""


def plot_activations(A,number_rows= 1,name="",i=0):
    A=A[0,:,:,:].detach().numpy()
    n_activations=A.shape[0]

    A_min=A.min().item()
    A_max=A.max().item()
    fig, axes = plt.subplots(number_rows, n_activations//number_rows)
    fig.subplots_adjust(hspace = 0.4)

    for i,ax in enumerate(axes.flat):
        if i< n_activations:
            # Set the label for the sub-plot.
            ax.set_xlabel( "activation:{0}".format(i+1))

            # Plot the image.
            ax.imshow(A[i,:], vmin=A_min, vmax=A_max, cmap='seismic')
            ax.set_xticks([])
            ax.set_yticks([])
    plt.show()

"""
Explanation: <code>plot_activation</code>: plot out the activations of the Convolutional layers
End of explanation
"""


def show_data(data_sample):

    plt.imshow(data_sample[0].numpy().reshape(28,28),cmap='gray')
    #print(data_sample[1].item())
    plt.title('y= '+ str(data_sample[1].item()))

"""
Explanation: <code>show_data</code>: plot out data sample
End of explanation
"""


train_dataset=dsets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
train_dataset

"""
Explanation: <a id="ref1"></a>
<h2 align=center>Prepare Data </h2>

Load the training dataset by setting the parameters <code>train </code> to <code>True</code>. The parameter <code>transform</code> defines the method to convert image to tensor.
End of explanation
"""


validation_dataset=dsets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
validation_dataset

"""
Explanation: Load the testing dataset by setting the parameters train  <code>False</code>. The parameter <code>transform</code> defines the method to convert image to tensor.
End of explanation
"""


train_dataset[0][1].type()

"""
Explanation: we can see the data type is long
End of explanation
"""


train_dataset[3][1]

"""
Explanation: Data Visualization
Each element in the rectangular  tensor corresponds to a number representing a pixel intensity  as demonstrated by  the following image.
<img src = "https://ibm.box.com/shared/static/7024mnculm8w2oh0080y71cpa48cib2k.png" width = 550, align = "center"></a>
Print out the third label
End of explanation
"""


show_data(train_dataset[3])

"""
Explanation: Plot the fourth sample
End of explanation
"""


show_data(train_dataset[2])

"""
Explanation: The fourth sample is a "1".
Plot the third sample.
End of explanation
"""


class CNN(nn.Module):
    def __init__(self,out_1=16,out_2=32):
        
        super(CNN,self).__init__()
        #first Convolutional layers
        self.cnn1=nn.Conv2d(in_channels=1,out_channels=out_1,kernel_size=5,padding=2)
        #activation function 
        self.relu1=nn.ReLU()
        #max pooling 
        self.maxpool1=nn.MaxPool2d(kernel_size=2)
        #second Convolutional layers
        self.cnn2=nn.Conv2d(in_channels=out_1,out_channels=out_2,kernel_size=5,stride=1,padding=2)
        #activation function 
        self.relu2=nn.ReLU()
        #max pooling 
        self.maxpool2=nn.MaxPool2d(kernel_size=2)
        #fully connected layer 
        self.fc1=nn.Linear(out_2*7*7,10)
        
    def forward(self,x):
        #first Convolutional layers
        out=self.cnn1(x)
        #activation function 
        out=self.relu1(out)
        #max pooling 
        out=self.maxpool1(out)
        #first Convolutional layers
        out=self.cnn2(out)
        #activation function
        out=self.relu2(out)
        #max pooling
        out=self.maxpool2(out)
        #flatten output 
        out=out.view(out.size(0),-1)
        #fully connected layer
        out=self.fc1(out)
        
        return out
    
    def activations(self,x):
        #outputs activation this is not necessary just for fun 
        z1=self.cnn1(x)
        a1=self.relu1(z1)
        out=self.maxpool1(a1)
        
        z2=self.cnn2(out)
        a2=self.relu2(z2)
        out=self.maxpool2(a2)
        out=out.view(out.size(0),-1)
        return z1,a1,z2,a2,out

"""
Explanation: <a id="ref3"></a>
Build a Convolutional Neral Network Class
Build a Convolutional Network class with two Convolutional layers and one fully connected layer. Pre-determine the size of the final output matrix. The parameters in the constructor are the number of output channels for the first and second layer.
End of explanation
"""


model=CNN(out_1=16,out_2=32)

"""
Explanation: <a id="ref3"></a>
<h2> Define the Convolutional Neral Network Classifier , Criterion function, Optimizer and Train the  Model  </h2>

There are 16 output channels for the first layer, and 32 outputs channels for the second layer
End of explanation
"""


plot_parameters(model.state_dict()['cnn1.weight'],number_rows=4,name="1st layer kernels before training ")
plot_parameters(model.state_dict()['cnn2.weight'],number_rows=4,name='2nd layer kernels before training' )

"""
Explanation: Plot the model parameters for the kernels before training the kernels. The kernels are initialized randomly.
End of explanation
"""


criterion=nn.CrossEntropyLoss()

"""
Explanation: Define the loss function
End of explanation
"""


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

"""
Explanation: Define the optimizer class
End of explanation
"""



train_loader=torch.utils.data.DataLoader(dataset=train_dataset,batch_size=100)
validation_loader=torch.utils.data.DataLoader(dataset=validation_dataset,batch_size=5000)

"""
Explanation: Define the dataset loader
End of explanation
"""


n_epochs=10
loss_list=[]
accuracy_list=[]
N_test=len(validation_dataset)
#n_epochs
for epoch in range(n_epochs):
    for x, y in train_loader:
      

        #clear gradient 
        optimizer.zero_grad()
        #make a prediction 
        z=model(x)
        # calculate loss 
        loss=criterion(z,y)
        # calculate gradients of parameters 
        loss.backward()
        # update parameters 
        optimizer.step()
        
        
        
    correct=0
    #perform a prediction on the validation  data  
    for x_test, y_test in validation_loader:

        z=model(x_test)
        _,yhat=torch.max(z.data,1)

        correct+=(yhat==y_test).sum().item()
        
   
    accuracy=correct/N_test

    accuracy_list.append(accuracy)
    
    loss_list.append(loss.data)


"""
Explanation: Train the model and determine validation accuracy technically test accuracy (This may take a long time)
End of explanation
"""


fig, ax1 = plt.subplots()
color = 'tab:red'
ax1.plot(loss_list,color=color)
ax1.set_xlabel('epoch',color=color)
ax1.set_ylabel('total loss',color=color)
ax1.tick_params(axis='y', color=color)
    
ax2 = ax1.twinx()  
color = 'tab:blue'
ax2.set_ylabel('accuracy', color=color)  
ax2.plot( accuracy_list, color=color)
ax2.tick_params(axis='y', labelcolor=color)
fig.tight_layout()

"""
Explanation: <a id="ref3"></a>
<h2 align=center>Analyse Results</h2>

Plot the loss and accuracy on the validation data:
End of explanation
"""


plot_channels(model.state_dict()['cnn1.weight'])
plot_channels(model.state_dict()['cnn2.weight'])

"""
Explanation: View the results of the parameters for the Convolutional layers
End of explanation
"""


show_data(train_dataset[1])

"""
Explanation: Consider the following sample
End of explanation
"""


out=model.activations(train_dataset[1][0].view(1,1,28,28))

A=out[1]
A.shape

"""
Explanation: Determine the activations
End of explanation
"""


plot_activations(out[0],number_rows=4,name="first feature map")

"""
Explanation: Plot out the first set of activations
End of explanation
"""


plot_activations(out[1],number_rows=4,name="first feature map after relu")

"""
Explanation: The image below is the result after applying the relu activation function
End of explanation
"""


plot_activations(out[2],number_rows=32//4,name="second feature map")

"""
Explanation: The image below is the result of the activation map after the second output layer.
End of explanation
"""


plot_activations(out[3],number_rows=4,name="second feature map after relu")

"""
Explanation: The image below is the result of the activation map after applying the second relu
End of explanation
"""


show_data(train_dataset[2])

out=model.activations(train_dataset[2][0].view(1,1,28,28))

plot_activations(out[0],number_rows=4,name="first feature map")

plot_activations(out[1],number_rows=4,name="first feature map ")

plot_activations(out[2],number_rows=4,name="second feature map ")

plot_activations(out[3],number_rows=4,name="second feature map ")

"""
Explanation: we can  see the result for the third sample
End of explanation
"""


count=0
for x,y in torch.utils.data.DataLoader(dataset=validation_dataset,batch_size=1):

    z=model(x)
    _,yhat=torch.max(z,1)
    if yhat!=y:
        show_data((x,y))

        plt.show()
        print("yhat:",yhat)
        count+=1
    if count>=5:
        break 
        

"""
Explanation: Plot the first five miss-classified samples:
End of explanation
"""


