#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct 21 13:51:08 2020

@author: ajv012
"""
# numpy and matplotlib
from numpy import asarray, vstack, savez_compressed, load, zeros, ones 
from numpy.random import randint
from matplotlib import pyplot
from os import listdir

# keras
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras.models import Model
from tensorflow.keras import Input
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import Conv2DTranspose, LeakyReLU, Activation, Concatenate, Dropout, BatchNormalization, LeakyReLU

# import helper functions 
from HelperFunctions import *
from Discriminator import *

# Define the generator network\

# Helper methods
# generate encoder blocks
def define_encoder_block(layer_in, n_filters, batchnorm=True):
  # initialize weights for the layer
  init = RandomNormal(stddev=0.02)

  # first downsampling layer
  g = Conv2D(n_filters, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(layer_in)
  
  # conditionally add batch normalization
  if batchnorm:
    g = BatchNormalization()(g, training=True)
	
  # leaky relu activation
  g = LeakyReLU(alpha=0.2)(g)
  return g

# generate decoder blocks 
def decoder_block(layer_in, skip_in, n_filters, dropout=True):

  # weight initialization
  init = RandomNormal(stddev=0.02)

  # add upsampling layer
  g = Conv2DTranspose(n_filters, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(layer_in)

  # add batch normalization
  g = BatchNormalization()(g, training=True)

  # conditionally add dropout
  if dropout:
    g = Dropout(0.5)(g, training=True)

  # merge with skip connection
  g = Concatenate()([g, skip_in])

  # relu activation
  g = Activation('relu')(g)
  return g

# define the standalone generator model
def define_generator(image_shape=(256,256,3)):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# image input
	in_image = Input(shape=image_shape)
	# encoder model
	e1 = define_encoder_block(in_image, 64, batchnorm=False)
	e2 = define_encoder_block(e1, 128)
	e3 = define_encoder_block(e2, 256)
	e4 = define_encoder_block(e3, 512)
	e5 = define_encoder_block(e4, 512)
	e6 = define_encoder_block(e5, 512)
	e7 = define_encoder_block(e6, 512)
	# bottleneck, no batch norm and relu
	b = Conv2D(512, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(e7)
	b = Activation('relu')(b)
	# decoder model
	d1 = decoder_block(b, e7, 512)
	d2 = decoder_block(d1, e6, 512)
	d3 = decoder_block(d2, e5, 512)
	d4 = decoder_block(d3, e4, 512, dropout=False)
	d5 = decoder_block(d4, e3, 256, dropout=False)
	d6 = decoder_block(d5, e2, 128, dropout=False)
	d7 = decoder_block(d6, e1, 64, dropout=False)
	# output
	g = Conv2DTranspose(3, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d7)
	out_image = Activation('tanh')(g)
	# define model
	model = Model(in_image, out_image)
	return model