#!/usr/bin/env python
# coding: utf-8

# # Q1.a.  (10 points) create an network to approximate sin(x)
# 
# In this problem, you will try to build a three-layer fully connected network to approximate the sin function using PyTorch. There are 100 hidden units in each layer and please use ReLU activation after the first two layers and no activation after the last layer. You should implement your network model as a subclass of [torch.nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) 

# ### <font color='red'>Define your network here</font>

# In[1]:


import torch
from torch import nn
import numpy as np
import random

class Mynetwork(torch.nn.Module):
    def __init__(self):
#-- complete the code here    


#-- end complete the code here    

    def forward(self,x):
#-- complete the code here    


#-- end complete the code here    
    


# ### Training

# In[2]:


model = Mynetwork()

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(10000):
    x=10. * torch.rand((200,1)).type(torch.FloatTensor)
    y=torch.sin(x).view(-1)
    
    optimizer.zero_grad()
    yh = model(x).view(-1)
    loss = loss_fn(yh,y)
    loss.backward()
    optimizer.step()

    if epoch % 1000 == 999:
        print(loss.item())


# ### Visualization

# In[3]:


from matplotlib import pyplot as plt

xv=np.array(range(100))/50. *np.pi
yv=model(torch.from_numpy(xv).type(torch.FloatTensor).view((100,-1))).view(-1).detach().numpy()

plt.plot(xv,yv,label='network')
plt.plot(xv,np.sin(xv),label='sin(x)')
plt.legend()


# # Q1.b. (Extra credit: 5 points) Repeat 1.a. using [torch.nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html) 
# 
# ### <font color='red'>Redefine your model below</font>

# In[4]:


# Hint: model = torch.nn.Sequential(...)

#-- complete the code here    



#-- end complete the code here    


# ### Training

# In[5]:


loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(10000):
    x=10. * torch.rand((200,1)).type(torch.FloatTensor)
    y=torch.sin(x).view(-1)
    
    optimizer.zero_grad()
    yh = model(x).view(-1)
    loss = loss_fn(yh,y)
    loss.backward()
    optimizer.step()

    if epoch % 1000 == 999:
        print(loss.item())


# ### Visualization

# In[6]:


from matplotlib import pyplot as plt

xv=np.array(range(100))/50. *np.pi
yv=model(torch.from_numpy(xv).type(torch.FloatTensor).view((100,-1))).view(-1).detach().numpy()

plt.plot(xv,yv,label='network')
plt.plot(xv,np.sin(xv),label='sin(x)')
plt.legend()


# # Q2. Classify clothing
# 
# In this problem, you will implement a network to classify different types of clothing. We will play with the [FachionMNIST](https://github.com/zalandoresearch/fashion-mnist) dataset, which contains ten different classes.
# 
# Your network should contain three conv layers with 6, 16, and 32 filters (all 3x3 and no padding), respectively. You should use ReLU activation after each conv layer followed by 2x2 max pooling. There will be two fully connected layers after the last pooling layer. The first fully connected layer has 64 neurons.

# # Q2.a. (2 points) How many outputs for the last fully connected layer (Hint: there are 10 classes)

# <font color='red'>Answer: </font>

# # Q2.b (6 points) How many parameters in the entire network? 

# <font color='red'>Answer: </font>
#     
# conv 1:
# 
# pooling 1:
# 
# conv2:
# 
# pooling 2:
# 
# conv3:
# 
# pooling 3:
# 
# fully-connected 1:
# 
# fully-connected 2:
# 
# 

# # Q2.c (12 points) Complete the code below by defining the network

# ### Download and load the dataset

# In[7]:


import torch
import torchvision
import torchvision.transforms as transforms

use_cuda = False # set it to True if you have GPU and want to use it

transform = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))])

# Create datasets for training & validation, download if necessary
training_set = torchvision.datasets.FashionMNIST('./data', train=True, transform=transform, download=True)
validation_set = torchvision.datasets.FashionMNIST('./data', train=False, transform=transform, download=True)

# Create data loaders for our datasets; shuffle for training, not for validation
training_loader = torch.utils.data.DataLoader(training_set, batch_size=4, shuffle=True, num_workers=2)
validation_loader = torch.utils.data.DataLoader(validation_set, batch_size=4, shuffle=False, num_workers=2)

# Class labels
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
        'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')

# Report split sizes
print('Training set has {} instances'.format(len(training_set)))
print('Validation set has {} instances'.format(len(validation_set)))


# ### Visualize the dataset

# In[8]:


import matplotlib.pyplot as plt
import numpy as np

# Helper function for inline image display
def matplotlib_imshow(img, one_channel=False):
    if one_channel:
        img = img.mean(dim=0)
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    if one_channel:
        plt.imshow(npimg, cmap="Greys")
    else:
        plt.imshow(np.transpose(npimg, (1, 2, 0)))

dataiter = iter(training_loader)
images, labels = dataiter.next()

# Create a grid from the images and show them
img_grid = torchvision.utils.make_grid(images)
matplotlib_imshow(img_grid, one_channel=True)
print('  '.join(classes[labels[j]] for j in range(4)))


# ### <font color='red'>Define your network here</font>

# In[9]:


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

class GarmentClassifier(nn.Module):
    def __init__(self):
#-- complete the code here    




#-- end complete the code here    

    def forward(self, x):
#-- complete the code here    




#-- end complete the code here    


if torch.cuda.is_available() and use_cuda:
    model = GarmentClassifier().cuda()
else:
    model = GarmentClassifier()


# ### Prepare "trainer"

# In[13]:


loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

def train_one_epoch(epoch_index):
    running_loss = 0.
    last_loss = 0.
    
    # Here, we use enumerate(training_loader) instead of
    # iter(training_loader) so that we can track the batch
    # index and do some intra-epoch reporting
    for i, data in enumerate(training_loader):
        # Every data instance is an input + label pair
        inputs, labels = data
        if torch.cuda.is_available() and use_cuda:
            inputs=inputs.cuda()
            labels=labels.cuda()

        # Zero your gradients for every batch!
        optimizer.zero_grad()

        # Make predictions for this batch
        outputs = model(inputs)

        # Compute the loss and its gradients
        loss = loss_fn(outputs, labels)
        loss.backward()

        # Adjust learning weights
        optimizer.step()

        # Gather data and report
        running_loss += loss.item()
        if i % 1000 == 999:
            last_loss = running_loss / 1000 # loss per batch
            print('  batch {} loss: {}'.format(i + 1, last_loss))
            running_loss = 0.

    return last_loss


# ### Start training

# In[15]:


# Initializing in a separate cell so we can easily add more epochs to the same run
torch.multiprocessing.set_sharing_strategy('file_system')
# timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
# writer = SummaryWriter('runs/fashion_trainer_{}'.format(timestamp))
epoch_number = 0

EPOCHS = 3

best_vloss = 1_000_000.

for epoch in range(EPOCHS):
    print('EPOCH {}:'.format(epoch_number + 1))

    # Make sure gradient tracking is on, and do a pass over the data
    model.train(True)
    avg_loss = train_one_epoch(epoch_number)

    # We don't need gradients on to do reporting
    model.train(False)

    running_vloss = 0.0
    for i, vdata in enumerate(validation_loader):
        vinputs, vlabels = vdata
        if torch.cuda.is_available() and use_cuda:
            vinputs=vinputs.cuda()
            vlabels=vlabels.cuda()
        voutputs = model(vinputs)
        vloss = loss_fn(voutputs, vlabels)
        running_vloss += vloss

    avg_vloss = running_vloss / (i + 1)
    print('LOSS train {} valid {}'.format(avg_loss, avg_vloss))

    # Track best performance, and save the model's state
    if avg_vloss < best_vloss:
        best_vloss = avg_vloss

    epoch_number += 1


# ### Verify results

# In[16]:


dataiter = iter(validation_loader)
images, labels = dataiter.next()

# Create a grid from the images and show them
img_grid = torchvision.utils.make_grid(images)
matplotlib_imshow(img_grid, one_channel=True)
print('Groundtruth: '+'  '.join(classes[labels[j]] for j in range(4)))
if torch.cuda.is_available() and use_cuda:
    images=images.cuda()
print('Result:      '+'  '.join([classes[i] for i in torch.argmax(model(images).cpu(),1)]))


# # Q2.d (Extra credit: 10 points) Try to see if you can modify the network so that the validation loss is less than 0.3 and the number of parameters are no more than 5% more than the original one 
# 

# # Q3. Object detection with YOLOv5
# 
# Please follow the instructions [here](https://pytorch.org/hub/ultralytics_yolov5/) to install YOLOv5 on PyTorch. 
# 
# 

# # Q3.a (10 points) Test YOLOv5s on one of your own photos. Show your code and results.

# # Q3.b (Extra credit: 5 points) Combine YOLO and OpenCV to detect object from webcam stream by completing the code below (Hint: dir(model(frame)))

# ### <font color='red'>Complete the code below</font>

# In[17]:


import cv2
import platform

cap=cv2.VideoCapture(0)

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

while (True):
    
    ret,frame=cap.read()
    
    results = model(frame)
    
#-- complete the code here


#-- end complete the code here    
    
    if cv2.waitKey(1) &0xFF == ord('q'): # press q or ESC to quit. You probably need to hit the screen first
        break
        
cap.release()
cv2.destroyAllWindows()


# In[ ]: