# -*- coding: utf-8 -*-
"""ReNS experiments - CIFAR10

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1byZ4xTfCK2x1Rhkxpl-Vv4sqA-bo4bis

# SETUP
"""

#@title Insatlling Pyorch

# !pip install torch
# !pip install torchvision

#@title Import Dependencies
import numpy as np
import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
from tqdm import tqdm
from typing import Optional, Union, Tuple, List, Sequence, Iterable
import math
from scipy.spatial.distance import euclidean
from torch.nn.modules.utils import _pair
from torchvision import models
from sklearn.metrics import jaccard_score
import matplotlib.pyplot as plt
from models.models import RegularAutoEncoder, ModulatedAutoEncoder, PseudoRecAutoEncoder


"""# TRAINING"""

batch_size = 32
num_epochs = 5

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

# Load MNIST data.
train_data = dsets.CIFAR10(root = './data', train = True,
                        transform = transform, download = True)

test_data = dsets.CIFAR10(root = './data', train = False,
                       transform = transform)

train_gen = torch.utils.data.DataLoader(dataset = train_data,
                                             batch_size = batch_size,
                                             shuffle = True)

test_gen = torch.utils.data.DataLoader(dataset = test_data,
                                      batch_size = batch_size,
                                      shuffle = False)

reflexor_size = 500
image_size = 32
channels = 3

# net = recurrentLayer(784, 784, 10, 5, 10, 0)
net1 = RegularAutoEncoder(channels * image_size ** 2, channels * image_size ** 2, reflexor_size)
net2 = ModulatedAutoEncoder(channels * image_size ** 2, channels * image_size ** 2, reflexor_size)
net3 = PseudoRecAutoEncoder(channels * image_size ** 2, channels * image_size ** 2, reflexor_size)



lr = .0001 # size of step
loss_function = nn.MSELoss()

# Unnormalize the image to display it
def img_fix(img):
  return np.transpose((img / 2 + 0.5).numpy(), (1, 2, 0))

# Commented out IPython magic to ensure Python compatibility.
train_losses = [[],[],[]]
test_losses = [[],[],[]]

real_imgs = [[],[],[]]
reconstructed_imgs = [[],[],[]]

param_counts = np.ones(3)

steps = [[],[],[]]

for num, net in enumerate([net1, net2, net3]):

  optimizer = torch.optim.Adam( net.parameters(), lr=lr)
  param_counts[num] = (sum(p.numel() for p in net.parameters() if p.requires_grad))

  for epoch in range(num_epochs):
    for i ,(images,labels) in enumerate(train_gen):
      #images = Variable(images.view(-1,28*28))
      labels = Variable(images.view(-1,3 * image_size ** 2))

      optimizer.zero_grad()
      outputs = net(images)
      loss = loss_function(outputs, labels)
      loss.backward()
      optimizer.step()

      if (i+1) % 300 == 0:
        temp_loss = loss.item()
        print('Epoch [%d/%d], Step [%d/%d], Loss: %.4f'
                   %(epoch+1, num_epochs, i+1, len(train_data)//batch_size, temp_loss))
        dupe = Variable(outputs[0].data, requires_grad=False)
        # plt.imshow(img_fix(images[0]))
        # plt.show()
        # plt.imshow(img_fix(dupe.view(3, image_size, image_size)))
        # plt.show()
        train_losses[num].append(temp_loss)
        steps[num].append((50000 * epoch) + ((i + 1) * batch_size))

        real_imgs[num].append(img_fix(images[0]))
        reconstructed_imgs[num].append(img_fix(dupe.view(3, image_size, image_size)))

        # Test Data

        score = 0
        total = 0
        for images,labels in test_gen:
          #images = Variable(images.view(-1,784))

          output = net(images)
          score += loss_function(output, images.view(-1, 3 * image_size ** 2)).item()
        test_losses[num].append((score))

plt.plot(steps[0], train_losses[0], label= "Baseline")
plt.plot(steps[1], train_losses[1], label= "Modulated")
plt.plot(steps[2], train_losses[2], label= "Recurrent with Modulation")
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Training loss history')
plt.legend()
plt.show()

plt.plot(steps[0], test_losses[0], label= "Baseline")
plt.plot(steps[1], test_losses[1], label= "Modulated")
plt.plot(steps[2], test_losses[2], label= "Recurrent with Modulation")
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Testing loss history')
plt.legend()
plt.show()

for num,count in enumerate(param_counts):
  param_counts[num] /= 1000


plt.bar(["Base", "Modulated", "ReNS"], param_counts)
plt.xlabel('Model')
plt.ylabel('# of thousands of Parameters')
plt.show()

from mpl_toolkits.axes_grid1 import ImageGrid

num_smaples = len(real_imgs[0])


for num in [0,1,2]:
  fig = plt.figure(figsize=(20.,20.))
  grid = ImageGrid(fig, 111,  # similar to subplot(111)
                  nrows_ncols=(2, num_smaples),  # creates 2x2 grid of axes
                  axes_pad=0.1,  # pad between axes in inch.
                  )

  for ax, im in zip(grid, real_imgs[num]+reconstructed_imgs[num]):
      # Iterating over the grid returns the Axes.
      ax.imshow(im)
      ax.axis("off")

  plt.show()
