# -*- coding: utf-8 -*-
"""

Automatically generated by Colaboratory.


"""

'''
Dec 17, 2020 Esha Singh
'''

import time
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms


if torch.cuda.is_available():
    torch.backends.cudnn.deterministic = True

# class ConvolutionalAutoencoder(torch.nn.Module):

#     def __init__(self):
#         super(ConvolutionalAutoencoder, self).__init__()
        
        
#         ### ENCODER
        
#         # 28x28x1 => 14x14x4  
#         self.conv_1 = torch.nn.Conv2d(in_channels=1,
#                                       out_channels=4,
#                                       kernel_size=(3, 3),
#                                       stride=(2, 2),
#                                       # floor((2(14-1) - 28 + 3) / 2) = 0
#                                       padding=0)
        
#         # 14x14x4 => 7x7x8
#         self.conv_2 = torch.nn.Conv2d(in_channels=4,
#                                       out_channels=8,
#                                       kernel_size=(3, 3),
#                                       stride=(2, 2),
#                                       # ceil((2(7-1) - 14 + 3) / 2) = 1
#                                       padding=1)                 
        
#         ### DECODER
                                         
#         # 7x7x8 => 15x15x4                          
#         self.deconv_1 = torch.nn.ConvTranspose2d(in_channels=8,
#                                                  out_channels=4,
#                                                  kernel_size=(3, 3),
#                                                  stride=(2, 2),
#                                                  padding=0)
        
#         # 15x15x4  => 29x29x1                           
#         self.deconv_2 = torch.nn.ConvTranspose2d(in_channels=4,
#                                                  out_channels=1,
#                                                  kernel_size=(3, 3),
#                                                  stride=(2, 2),
#                                                  padding=1)
#     def forward(self, x):
        
#         ### ENCODER
#         x = self.conv_1(x)
#         x = F.leaky_relu(x)
#         x = self.conv_2(x)
#         x = F.leaky_relu(x)
          
#         ### DECODER
#         x = self.deconv_1(x)  
#         x = F.leaky_relu(x)
#         x = self.deconv_2(x)
#         x = F.leaky_relu(x)
#         x = x[:, :, :-1, :-1]
#         x = torch.sigmoid(x)
#         return x

    
# torch.manual_seed(random_seed)
# model = ConvolutionalAutoencoder()
# model = model.to(device)

# optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

### MODEL TO BE USED FOR DELG

random_seed = 0
learning_rate = 0.001
num_epochs = 50
batch_size = 128

# Architecture
num_classes = 10
num_features = 784
num_latent = 50


def to_onehot(labels, num_classes, device):

    labels_onehot = torch.zeros(labels.size()[0], num_classes).to(device)
    labels_onehot.scatter_(1, labels.view(-1, 1), 1)

    return labels_onehot


class ConditionalVariationalAutoencoder(torch.nn.Module):

    def __init__(self, num_features, num_latent, num_classes):
        super(ConditionalVariationalAutoencoder, self).__init__()
        
        self.num_classes = num_classes
        
        # ENCODER
        ##############

        self.enc_conv_1 = torch.nn.Conv2d(in_channels=1+self.num_classes,
                                          out_channels=16,
                                          kernel_size=(6, 6),
                                          stride=(2, 2),
                                          padding=0)

        self.enc_conv_2 = torch.nn.Conv2d(in_channels=16,
                                          out_channels=32,
                                          kernel_size=(4, 4),
                                          stride=(2, 2),
                                          padding=0)                 
        
        self.enc_conv_3 = torch.nn.Conv2d(in_channels=32,
                                          out_channels=64,
                                          kernel_size=(2, 2),
                                          stride=(2, 2),
                                          padding=0)                     
        
        self.z_mean = torch.nn.Linear(64*2*2, num_latent)
        self.z_log_var = torch.nn.Linear(64*2*2, num_latent)



               
        # DECODER
        self.dec_linear_1 = torch.nn.Linear(num_latent+self.num_classes, 64*2*2)
        self.dec_deconv_1 = torch.nn.ConvTranspose2d(in_channels=64,
                                                     out_channels=32,
                                                     kernel_size=(2, 2),
                                                     stride=(2, 2),
                                                     padding=0)
                                 
        self.dec_deconv_2 = torch.nn.ConvTranspose2d(in_channels=32,
                                                     out_channels=16,
                                                     kernel_size=(4, 4),
                                                     stride=(3, 3),
                                                     padding=1)
        
        self.dec_deconv_3 = torch.nn.ConvTranspose2d(in_channels=16,
                                                     out_channels=11,
                                                     kernel_size=(6, 6),
                                                     stride=(3, 3),
                                                     padding=4)        


    def reparameterize(self, z_mu, z_log_var):
        # Sample epsilon from standard normal distribution
        eps = torch.randn(z_mu.size(0), z_mu.size(1)).to(device)
        z = z_mu + eps * torch.exp(z_log_var/2.) 
        return z

    def encoder(self, features, targets):
        onehot_targets = to_onehot(targets, self.num_classes, device)
        onehot_targets = onehot_targets.view(-1, self.num_classes, 1, 1)
        
        ones = torch.ones(features.size()[0], 
                          self.num_classes,
                          features.size()[2], 
                          features.size()[3], 
                          dtype=features.dtype).to(device)
        ones = ones * onehot_targets
        x = torch.cat((features, ones), dim=1)
        
        x = self.enc_conv_1(x)
        x = F.leaky_relu(x)
       
        x = self.enc_conv_2(x)
        x = F.leaky_relu(x)
       
        x = self.enc_conv_3(x)
        x = F.leaky_relu(x)
        
        z_mean = self.z_mean(x.view(-1, 64*2*2))
        z_log_var = self.z_log_var(x.view(-1, 64*2*2))
        encoded = self.reparameterize(z_mean, z_log_var)
        return z_mean, z_log_var, encoded
    

    def decoder(self, encoded, targets):
        onehot_targets = to_onehot(targets, self.num_classes, device)
        encoded = torch.cat((encoded, onehot_targets), dim=1)        
        
        x = self.dec_linear_1(encoded)
        x = x.view(-1, 64, 2, 2)
        
        x = self.dec_deconv_1(x)
        x = F.leaky_relu(x)
        #print('deconv1 out:', x.size())
        
        x = self.dec_deconv_2(x)
        x = F.leaky_relu(x)
        
        x = self.dec_deconv_3(x)
        x = F.leaky_relu(x)
        #print('deconv1 out:', x.size())
        
        decoded = torch.sigmoid(x)
        return decoded

    def forward(self, features, targets):
        z_mean, z_log_var, encoded = self.encoder(features, targets)
        decoded = self.decoder(encoded, targets)
        return z_mean, z_log_var, encoded, decoded

    
torch.manual_seed(random_seed)
model = ConditionalVariationalAutoencoder(num_features,
                                          num_latent,
                                          num_classes)
model = model.to(device)
    



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

# Commented out IPython magic to ensure Python compatibility.

for epoch in range(num_epochs):
    for batch_idx, (features, targets) in enumerate(train_loader):
        
        features = features.to(device)
        targets = targets.to(device)

        z_mean, z_log_var, encoded, decoded = model(features, targets)

        kl_divergence = (0.5 * (z_mean**2 + 
                                torch.exp(z_log_var) - z_log_var - 1)).sum()
        
        
        onehot_targets = to_onehot(targets, num_classes, device)
        onehot_targets = onehot_targets.view(-1, num_classes, 1, 1)
        
        ones = torch.ones(features.size()[0], 
                          num_classes,
                          features.size()[2], 
                          features.size()[3], 
                          dtype=features.dtype).to(device)
        ones = ones * onehot_targets
        x_con = torch.cat((features, ones), dim=1)
        
        
        ### Compute loss
        pixelwise_bce = F.binary_cross_entropy(decoded, x_con, reduction='sum')
        cost = kl_divergence + pixelwise_bce
        
        optimizer.zero_grad()
        cost.backward()
        optimizer.step()
        
        if not batch_idx % 50:
            print ('Epoch: %03d/%03d | Batch %03d/%03d | Cost: %.4f' 
#                    %(epoch+1, num_epochs, batch_idx, 
                     len(train_loader), cost))
            
    print('Time elapsed: %.2f min' % ((time.time() - start_time)/60))
    
print('Total Training Time: %.2f min' % ((time.time() - start_time)/60))

"""Evaluation - Reconstruction"""

matplotlib inline
import matplotlib.pyplot as plt

### VISUALIZATION

n_images = 15
image_width = 28

fig, axes = plt.subplots(nrows=2, ncols=n_images, 
                         sharex=True, sharey=True, figsize=(20, 2.5))
orig_images = features[:n_images]
decoded_images = decoded[:n_images, 0]

for i in range(n_images):
    for ax, img in zip(axes, [orig_images, decoded_images]):
        ax[i].imshow(img[i].detach().to(torch.device('cpu')).reshape((image_width, image_width)), cmap='binary')