import torch 
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pylab as plt
import numpy as np
torch.manual_seed(0)

"""
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>Plot Output of Neural Networks with One Hidden Layer and Activations for Each Iteration of Gradient Descent</font></h1> 

# Table of Contents
In this lab, you will use a single-layer neural network 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">Neural Network Module and Training Function</a></li>
<li><a href="#ref2">Prepare Data</a></li>
<li><a href="#ref3">Define the Neural Network, Criterion Function, Optimizer, and Train the  Model</a></li>

<li><a href="#ref4">Practice Question</a></li>
<br>
<p></p>
Estimated Time Needed: <strong>25 min</strong>
</div>

<hr>

Import the following libraries:
End of explanation
"""


def PlotStuff(X,Y,model,epoch,leg=True):
    plt.plot(X.numpy(),model(X).detach().numpy(),label='epoch '+str(epoch))
    plt.plot(X.numpy(),Y.numpy(),'r')
    plt.xlabel('x')
    if leg==True:
        plt.legend()
    else:
        pass

"""
Explanation: <a id="ref0"></a>
<h2 align=center>Helper Functions</h2>
End of explanation
"""


class Net(nn.Module):
    def __init__(self,D_in,H,D_out):
        super(Net,self).__init__()
        self.linear1=nn.Linear(D_in,H)
        self.linear2=nn.Linear(H,D_out)
        #define the first linear layer as an attribute, this is not good practice
        self.a1=None
        self.l1=None
        
    def forward(self,x):
        self.l1=self.linear1(x)
        self.a1=F.sigmoid(self.l1)
        yhat=F.sigmoid(self.linear2(self.a1))
        return yhat

"""
Explanation: <a id="ref1"></a>
<h2 align=center>Neural Network Module and Training Function </h2>

Define the activations and the output of the first linear layer as an attribute. Note that this is not good practice.
End of explanation
"""


def train(Y,X,model,optimizer,criterion,epochs=1000):
    error=[]

    for epoch in range(epochs):

        for y,x in zip(Y,X):
            yhat=model(x)
            loss=criterion(yhat,y)
            loss.backward()
            optimizer.step()
           #clear gradient 
            optimizer.zero_grad()
        error.append(loss.item())
        if epoch % 300==0:    
            PlotStuff(X,Y,model,epoch,leg=True)
            plt.show()
            model(X)
            plt.scatter( model.a1.detach().numpy()[:,0], model.a1.detach().numpy()[:,1],c=Y.numpy().reshape(-1))
            plt.title('activations')
            plt.show()
    return  error

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


X=torch.arange(-20,20,1).view(-1,1).type(torch.FloatTensor )
Y=torch.zeros(X.shape[0])
Y[(X[:,0]>-4)& (X[:,0]<4)]=1.0

"""
Explanation: <a id="ref2"></a>
<h2 align=center>Prepare Data </h2>
End of explanation
"""


def criterion_cross(outputs,labels):
    out=-1*torch.mean(labels*torch.log(outputs) +(1-labels)*torch.log(1-outputs))
    return out

"""
Explanation: <a id="ref3"></a>
<h2 align=center>Define the Neural Network, Criterion Function, Optimizer. and Train the  Model</h2>

Create the Cross-Entropy loss function:
End of explanation
"""


D_in=1
H=2
D_out=1

learning_rate=0.1
model=Net(D_in,H,D_out)
optimizer=torch.optim.SGD(model.parameters(), lr=learning_rate)
cost_cross=train(Y,X,model,optimizer,criterion_cross,epochs=1000)
plt.plot(cost_cross)
plt.xlabel('epoch')
plt.title('cross entropy loss ')

"""
Explanation: Define the Neural Network, Optimizer, and Train the Model:
End of explanation
"""


learning_rate=0.1
criterion_mse=nn.MSELoss()
model=Net(D_in,H,D_out)
optimizer=torch.optim.SGD(model.parameters(), lr=learning_rate)
cost_mse=train(Y,X,model,optimizer,criterion_mse,epochs=1000)
plt.plot(cost_mse)
plt.xlabel('epoch')
plt.title('MSE loss ')

"""
Explanation: By examining the output of the  activation, you see by the 600th epoch that the data has been mapped to a linearly separable space.
<a id="ref3"></a>
<h2 align=center>Practice Question </h2>

Repeat the previous steps above by using the MSE cost or total loss:
End of explanation
"""


