import torch
import numpy as np
import matplotlib.pyplot as plt
from torch import nn,optim
from mpl_toolkits.mplot3d import Axes3D
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms

"""
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>

<h1 align=center><font size = 5>Linear Regression Multiple Outputs</font></h1> 
# Table of Contents
In this lab, you will create a model the Pytroch way. This will help you as models get more complicated.

<div class="alert alert-block alert-info" style="margin-top: 20px">
<li><a href="#ref0">Make Some Data</a></li>
<li><a href="#ref1">Create the Model and Cost Function the Pytorch way</a></li>
<li><a href="#ref2">Train the Model: Batch Gradient Descent</a></li>
<li><a href="#ref3">Practice Questions </a></li>
<br>
<p></p>
Estimated Time Needed: <strong>20 min</strong>
</div>

<hr>

Import the following libraries:
End of explanation
"""


torch.manual_seed(1)

"""
Explanation: Set the random seed:
End of explanation
"""


from torch.utils.data import Dataset, DataLoader
class Data(Dataset):
    def __init__(self):
            self.x=torch.zeros(20,2)
            self.x[:,0]=torch.arange(-1,1,0.1)
            self.x[:,1]=torch.arange(-1,1,0.1)
            self.w=torch.tensor([ [1.0,-1.0],[1.0,3.0]])
            self.b=torch.tensor([[1.0,-1.0]])
            self.f=torch.mm(self.x,self.w)+self.b
            
            self.y=self.f+0.001*torch.randn((self.x.shape[0],1))
            self.len=self.x.shape[0]

    def __getitem__(self,index):

        return self.x[index],self.y[index]
    
    def __len__(self):
        return self.len

"""
Explanation: <a id="ref0"></a>
<h2 align=center>Make Some Data </h2>
Create a dataset class with two-dimensional features and two targets:
End of explanation
"""


data_set=Data()

"""
Explanation: create a dataset object
End of explanation
"""


class linear_regression(nn.Module):
    def __init__(self,input_size,output_size):
        super(linear_regression,self).__init__()
        self.linear=nn.Linear(input_size,output_size)
    def forward(self,x):
        yhat=self.linear(x)
        return yhat

"""
Explanation: <a id="ref1"></a>
<h2 align=center>Create the Model, Optimizer, and Total Loss Function (cost)</h2>

Create a custom module:
End of explanation
"""


model=linear_regression(2,2)

"""
Explanation: Create an optimizer object and set the learning rate to 0.1. Don't forget to enter the model parameters in the constructor.
End of explanation
"""


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

"""
Explanation: Create an optimizer object and set the learning rate to 0.1. Don't forget to enter the model parameters in the constructor.  
<img src = "https://ibm.box.com/shared/static/f8hskuwrnctjg21agud69ddla0jkbef5.png" width = 100, align = "center">
End of explanation
"""


criterion = nn.MSELoss()

"""
Explanation: Create the criterion function that calculates the total loss or cost:
End of explanation
"""


train_loader=DataLoader(dataset=data_set,batch_size=5)

"""
Explanation: Create a data loader object and set the batch_size to 5:
End of explanation
"""


LOSS=[]
 
epochs=100
   
for epoch in range(epochs):
    for x,y in train_loader:
        #make a prediction 
        yhat=model(x)
        #calculate the loss
        loss=criterion(yhat,y)
        #store loss/cost 
        LOSS.append(loss.item())
        #clear gradient 
        optimizer.zero_grad()
        #Backward pass: compute gradient of the loss with respect to all the learnable parameters
        loss.backward()
        #the step function on an Optimizer makes an update to its parameters
        optimizer.step()
     


    

"""
Explanation: <a id="ref2"></a>
<h2 align=center>Train the Model via Mini-Batch Gradient Descent </h2>

Run 100 epochs of Mini-Batch Gradient Descent and store the total loss or cost for every iteration. Remember that this is an approximation of the true total loss or cost.
End of explanation
"""


plt.plot(LOSS)
plt.xlabel("iterations ")
plt.ylabel("Cost/total loss ")

"""
Explanation: Plot the cost:
End of explanation
"""


