import os,time
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from keras import initializers
from keras.datasets import cifar10
from keras.initializers import RandomNormal
from keras.layers import BatchNormalization, Conv2D, Conv2DTranspose, Dense,Dropout, Flatten, Input, Reshape, UpSampling2D,ZeroPadding2D
from keras.layers.advanced_activations import LeakyReLU
from keras.models import Model, Sequential
from keras.optimizers import Adam
from PIL import Image, ImageDraw
np.random.seed(1337)

#function to create the generator model
def create_generator(d,noise_dim,optimizer):
  generator = Sequential()  
  generator.add(Dense(d*d*256,kernel_initializer=RandomNormal(0,0.02),input_dim=noise_dim))
  generator.add(LeakyReLU(0.2))
  generator.add(Reshape((d,d,256)))
  generator.add(Conv2DTranspose(128,(4,4),strides=2,padding='same',kernel_initializer=RandomNormal(0,0.02)))
  generator.add(LeakyReLU(0.2))  
  generator.add(Conv2DTranspose(128,(4,4),strides=2,padding='same',kernel_initializer=RandomNormal(0,0.02)))
  generator.add(LeakyReLU(0.2)) 
  generator.add(Conv2DTranspose(128,(4,4),strides=2,padding='same',kernel_initializer=RandomNormal(0,0.02)))
  generator.add(LeakyReLU(0.2))
  generator.add(Conv2D(3,(3,3),padding='same',activation='tanh',kernel_initializer=RandomNormal(0,0.02)))
  generator.compile(loss='binary_crossentropy', optimizer=optimizer)
  return generator

#function to create the discriminator module
def create_discriminator(w,optimizer):
  discriminator = Sequential()    
  discriminator.add(Conv2D(64,(3,3),padding='same',kernel_initializer=RandomNormal(0,0.02),input_shape=(w,w,3)))
  discriminator.add(LeakyReLU(0.2))  
  discriminator.add(Conv2D(128,(3,3),strides=2,padding='same',kernel_initializer=RandomNormal(0,0.02)))
  discriminator.add(LeakyReLU(0.2))
  discriminator.add(Conv2D(128,(3,3),strides=2,padding='same',kernel_initializer=RandomNormal(0,0.02)))
  discriminator.add(LeakyReLU(0.2))  
  discriminator.add(Conv2D(256,(3,3),strides=2,padding='same',kernel_initializer=RandomNormal(0,0.02)))
  discriminator.add(LeakyReLU(0.2))  
  discriminator.add(Flatten())
  discriminator.add(Dropout(0.4))
  discriminator.add(Dense(1,activation='sigmoid',input_shape=(w,w,3)))  
  discriminator.compile(loss='binary_crossentropy',optimizer=optimizer)
  return discriminator

#function to plot 5 images from the training dataset
def plot_img(train10X):
  for i in range(5):
    # define subplot
	  plt.subplot(7, 7, 1 + i)
	  # turn off axis
	  plt.axis('off')
	  # plot raw pixel data
	  plt.imshow(train10X[i])
  plt.show()

#function to display images generated by generator
def show_images(noise,epoch,generator,w):
  channels=3
  generated_images = generator.predict(noise)
  plt.figure(figsize=(10,10))
  for i, image in enumerate(generated_images):
    plt.subplot(10,10,i+1)
    if channels == 1:
      plt.imshow(np.clip(image.reshape((w,w)),0.0,1.0),cmap='gray')
    else:
      plt.imshow(np.clip(image.reshape((w,w,channels)),0.0,1.0))
    plt.axis('off')
    
  plt.tight_layout()

def gan(w,optimizer,noise_dim,train10X,d,epochs,steps_per_epoch,batch_size):
  #create the generator and discriminator module
  discriminator = create_discriminator(w,optimizer)
  generator = create_generator(d,noise_dim,optimizer)

  # Make the discriminator untrainable when we are training the generator
  discriminator.trainable = False

  #link the two models to create the GAN
  gan_input = Input(shape=(noise_dim,))
  fake_image = generator(gan_input)

  gan_output = discriminator(fake_image)

  gan = Model(gan_input, gan_output)
  gan.compile(loss='binary_crossentropy', optimizer=optimizer)

  #constant noise for viewing how the GAN progresses
  static_noise = np.random.normal(0, 1, size=(100, noise_dim))

  #training loop
  for epoch in range(epochs):
    for batch in range(steps_per_epoch):
      #add noise
      noise = np.random.normal(0, 1, size=(batch_size, noise_dim))
      #shuffle and get the real image from training set based on the batch szie every time
      real_x = train10X[np.random.randint(0, train10X.shape[0], size=batch_size)]
      #generate fake image
      fake_x = generator.predict(noise)
      x = np.concatenate((real_x, fake_x))
      disc_y = np.zeros(2*batch_size)
      disc_y[:batch_size] = 0.9
      #discriminator loss
      d_loss = discriminator.train_on_batch(x, disc_y)
      d_loss=np.round(d_loss,4) #limit output to 4 decimal places
      y_gen = np.ones(batch_size)
      #generator loss
      g_loss = gan.train_on_batch(noise, y_gen)
      g_loss=np.round(g_loss,4) #limit output to 4 decimal places
    im1 = tf.convert_to_tensor(real_x, tf.float32)
    im2 = tf.convert_to_tensor(fake_x, tf.float32)
    ssim_diff=tf.image.ssim(im1,im2,max_val=1, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03)
    avg_ssim=np.average(ssim_diff)
    avg_ssim="{:0.4f}".format(avg_ssim) #limit output to 4 decimal places
    print(f'Epoch: {epoch}') 
    print(f'SSIM Score: {ssim_diff}')
    print(f'Discriminator Loss: {d_loss} \t\t Generator Loss: {g_loss} \t\t Average SSIM score: {avg_ssim}')

    if epoch % 2 == 0:
        show_images(static_noise,epoch,generator,w)

def main():
  #size
  d=4
  # The dimension of noise
  noise_dim = 100
  batch_size = 10
  steps_per_epoch = 312
  epochs = 400
  w=32
  optimizer = Adam(0.0002, 0.5)
  #get cifar 10 data
  (train10X, train10y), (test10X, test10y) = cifar10.load_data()
  #normalize the data
  X_train_mean = np.mean(train10X, axis=(0,1,2))
  X_train_std = np.std(train10X, axis=(0,1,2))
  train10X = (train10X - X_train_mean) / X_train_std
  test10X = (test10X - X_train_mean) / X_train_std
  #one hot encoding
  Y_train10 = tf.one_hot(train10y,10)
  Y_test10 = tf.one_hot(test10y,10)
  # summarize loaded dataset
  print('Train Cifar10: X=%s, y=%s' % (train10X.shape, Y_train10.shape))
  print('Test Cifar10: X=%s, y=%s' % (test10X.shape, Y_test10.shape))
  #plot the training dataset 
  plot_img(train10X)
  #call gan function
  st=time.time()
  gan(w,optimizer,noise_dim,train10X,d,epochs,steps_per_epoch,batch_size)
  print(f'Total time required: {np.round(time.time()-st,4)}')

main()

