from keras.layers.core import Activation
from keras.layers.normalization import BatchNormalization
from keras.layers.convolutional import UpSampling2D
# from keras.layers.core import Flatten
from keras.layers import Input
# from keras.layers.convolutional import Conv2DTranspose
from keras.models import Model
from keras.layers.advanced_activations import LeakyReLU, PReLU
from keras.layers import add
from keras.layers import Lambda
from keras.layers import Maximum
from aPcon_keras_debug import aPConv2D as PConv2D

# Residual block
def res_block_gen(model, kernal_size, filters, strides):
    
    gen = model
    
    model = PConv2D(filters = filters, kernel_size = kernal_size, strides = strides, padding = "same", activation="relu")(model)
    model = BatchNormalization(momentum = 0.5)(model)
    model = PConv2D(filters = filters, kernel_size = kernal_size, strides = strides, padding = "same", activation="relu")(model)
    model = BatchNormalization(momentum = 0.5)(model)
        
    model = add([gen, model])
    
    return model
    
    
def up_sampling_block(model, kernal_size, filters, strides):
    
    model = PConv2D(filters = filters, kernel_size = kernal_size, strides = strides, padding = "same")(model)
    model = UpSampling2D(size = 2)(model)
    
    return model

class Generator(object):

    def __init__(self, noise_shape):
        
        self.noise_shape = noise_shape
        

    def generator(self, upsample):
        gen_input = Input(shape = self.noise_shape)
        model = PConv2D(filters = 128, kernel_size = 9, strides = 1, padding = "same", activation="relu")(gen_input)
        gen_model = model
        
        for index in range(4): #originally it was 16
	        model = res_block_gen(model, 3, 128, 1)
            
        model = PConv2D(filters = 128, kernel_size = 3, strides = 1, padding = "same", activation="relu")(model)
        model = BatchNormalization(momentum = 0.5)(model)
        model = add([gen_model, model])
        
	    # Using 2 UpSampling Blocks, or the number of updsampling provided 
        for index in range(upsample):
        # for index in range(2):
    	        model = up_sampling_block(model, 3, 256, 1)
        model = PConv2D(filters = 3, kernel_size = 9, strides = 1, padding = "same")(model) # filters=1 mean the ouput will be a one channel image.
        model = Activation('tanh')(model)
        generator_model = Model(inputs = gen_input, outputs = model)
        
        return generator_model
