#!/usr/bin/env python
from __future__ import print_function
import time
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np
from torch import nn, optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import pandas as pd
import sys, json
sys.path.append(".")
from generator import G, Generator
from discriminator import D, Discriminator



def preprocessing():
    with open('../train.json') as json_file:
        dict_train = json.load(json_file)

    id_ = []
    cuisine = []
    ingredients = []
    for i in range(len(dict_train)):
        id_.append(dict_train[i]['id'])
        cuisine.append(dict_train[i]['cuisine'])
        ingredients.append(dict_train[i]['ingredients'])
    df = pd.DataFrame({'id':id_, 
                       'cuisine':cuisine, 
                       'ingredients':ingredients})
    print(df.head(10))

    return df


## Parameters ## 
def set_params(netD, netG):
    EPOCH = 20 # play with me
    LR = 0.001
    criterion = nn.BCELoss()
    optimizerD = optim.Adam(netD.parameters(), lr=LR, betas=(0.5, 0.999))
    optimizerG = optim.Adam(netG.parameters(), lr=LR, betas=(0.5, 0.999))
    return EPOCH, LR, criterion, optimizerD, optimizerG

def weights_init(m):
    """
    Takes as input a neural network m that will initialize all its weights.
    """
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

def define_GAN():
    # Creating the Generator
    netG = G()
    netG.apply(weights_init)
    # Creating the discriminat
    netD = D()
    netD.apply(weights_init) 
    return netG, netD




def main():
    print("Working")
    

    dataloader = preprocessing()
    netG, netD = define_GAN()
    EPOCH, LR, criterion, optimizerD, optimizerG = set_params(netD, netG)

    print(dataloader.head(10))
    for epoch in range(EPOCH):
           for i, data in enumerate(dataloader, 0):
	       # 1st Step: Updating the weights of the neural network of the discriminator
               netD.zero_grad()
               
               # Training the discriminator with a real image of the dataset
               real,_ = data
               input = Variable(real)
               target = Variable(torch.ones(input.size()[0]))
               output = netD(input)
               errD_real = criterion(output, target)
               
               # Training the discriminator with a fake image generated by the generator
               noise = Variable(torch.randn(input.size()[0], 100, 1, 1))
               fake = netG(noise)
               target = Variable(torch.zeros(input.size()[0]))
               output = netD(fake.detach())
               errD_fake = criterion(output, target)
               
               # Backpropagating the total error
               errD = errD_real + errD_fake
               errD.backward()
               optimizerD.step()
               
               # 2nd Step: Updating the weights of the neural network of the generator
               netG.zero_grad()
               target = Variable(torch.ones(input.size()[0]))
               output = netD(fake)
               errG = criterion(output, target)
               errG.backward()
               optimizerG.step()
               
               # 3rd Step: Printing the losses and saving the real images and the generated images of the minibatch every 100 steps
               print('[%d/%d][%d/%d] Loss_D: %.4f; Loss_G: %.4f' % (epoch, EPOCH, i, len(dataloader), errD.item(), errG.item()))
               if i % 100 == 0:
                   vutils.save_image(real, '%s/real_samples.png' % "./results", normalize=True)
                   fake = netG(noise)
                   vutils.save_image(fake.data, '%s/fake_samples_epoch_%03d.png' % ("./results", epoch), normalize=True)


if __name__=="__main__":
    main()
