from keras.models import Model
from keras.layers import Input, Dense, Flatten, Dropout, Reshape, Conv2D, MaxPooling2D, UpSampling2D, Activation, Lambda
from keras.layers import Normalization
from keras import regularizers
import tensorflow as tf
from keras import backend as K
from keras.constraints import unit_norm, max_norm
from keras.losses import mean_absolute_error


# reparameterization trick
def sampling(args):
    """Reparameterization trick by sampling fr an isotropic unit Gaussian.
    # Arguments:
        args (tensor): mean and log of variance of Q(z|X)
    # Returns:
        z (tensor): sampled latent vector
    """
    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    # by default, random_normal has mean=0 and std=1.0
    epsilon = K.random_normal(shape=(batch, dim),dtype=tf.float64)
    return z_mean + K.exp(0.5 * z_log_var) * epsilon

def build_vae(input_shape, patch_x, patch_y, ft_bank_baseline, latent_dim, dropout_alpha, L2_reg, binary_image):
    # build encoder model
    input_r = Input(shape=(1,), name='ground_truth',dtype=tf.float64)
    input_image = Input(shape=input_shape, name='input_image',dtype=tf.float64)

    feature = Conv2D(ft_bank_baseline, activation='relu', kernel_size=(3, 3), padding='same',dtype=tf.float64)(input_image)
    feature = MaxPooling2D(pool_size=(2, 2),dtype=tf.float64)(feature)

    feature = Conv2D(ft_bank_baseline * 2, activation='relu', kernel_size=(3, 3), padding='same',dtype=tf.float64)(feature)
    feature = MaxPooling2D(pool_size=(2, 2),dtype=tf.float64)(feature)

    feature = Conv2D(ft_bank_baseline * 4, activation='relu', kernel_size=(3, 3), padding='same',dtype=tf.float64)(feature)
    feature = MaxPooling2D(pool_size=(2, 2),dtype=tf.float64)(feature)

    feature = Flatten(dtype=tf.float64)(feature)
    feature = Dropout(dropout_alpha,dtype=tf.float64)(feature)
    feature_dense = Dense(latent_dim * 4, activation='tanh', kernel_regularizer=regularizers.l2(L2_reg),dtype=tf.float64)(feature)

    feature_z_mean = Dense(latent_dim * 2, activation='tanh',dtype=tf.float64)(feature_dense)
    z_mean = Dense(latent_dim, name='z_mean',dtype=tf.float64)(feature_z_mean)
    feature_z_log_var = Dense(latent_dim * 2, activation='tanh',dtype=tf.float64)(feature_dense)
    z_log_var = Dense(latent_dim, name='z_log_var',dtype=tf.float64)(feature_z_log_var)

    feature_r_mean = Dense(latent_dim * 2, activation='tanh',dtype=tf.float64)(feature_dense)
    r_mean = Dense(1, name='r_mean',dtype=tf.float64)(feature_r_mean)
    feature_r_log_var = Dense(latent_dim * 2, activation='tanh',dtype=tf.float64)(feature_dense)
    r_log_var = Dense(1, name='r_log_var',dtype=tf.float64)(feature_r_log_var)

    # use reparameterization trick to push the sampling out as input
    z = Lambda(sampling, output_shape=(latent_dim,), name='z',dtype=tf.float64)([z_mean, z_log_var])
    r = Lambda(sampling, output_shape=(1,), name='r',dtype=tf.float64)([r_mean, r_log_var])

    # instantiate encoder model
    encoder = Model([input_image, input_r], [z_mean, z_log_var, z, r_mean, r_log_var, r], name='encoder')

    # build generator model
    generator_input = Input(shape=(1,), name='genrator_input',dtype=tf.float64)
    pz_mean = Dense(latent_dim, name='pz_mean', kernel_constraint=unit_norm(),dtype=tf.float64)(generator_input)
    pz_log_var = Dense(1, name='pz_log_var', kernel_constraint=max_norm(0),dtype=tf.float64)(generator_input)

    # instantiate generator model
    generator = Model(generator_input, [pz_mean, pz_log_var], name='generator')

    # build decoder model
    latent_input = Input(shape=(latent_dim,), name='z_sampling',dtype=tf.float64)
    decoded = Dense(latent_dim * 2, activation='tanh', kernel_regularizer=regularizers.l2(L2_reg),dtype=tf.float64)(latent_input)
    decoded = Dense(latent_dim * 4, activation='tanh', kernel_regularizer=regularizers.l2(L2_reg),dtype=tf.float64)(decoded)
    decoded = Dense(int(patch_x / 4 * patch_y / 4 * ft_bank_baseline * 4), activation='relu', kernel_regularizer=regularizers.l2(L2_reg),dtype=tf.float64)(decoded)
    decoded = Reshape((int(patch_x / 4), int(patch_y / 4), ft_bank_baseline * 4),dtype=tf.float64)(decoded)

    decoded = Conv2D(ft_bank_baseline * 4, kernel_size=(3, 3), padding='same',dtype=tf.float64)(decoded)
    decoded = Activation('relu',dtype=tf.float64)(decoded)
    decoded = UpSampling2D((2, 2),dtype=tf.float64)(decoded)

    decoded = Conv2D(ft_bank_baseline * 2, kernel_size=(3, 3), padding='same',dtype=tf.float64)(decoded)
    decoded = Activation('relu',dtype=tf.float64)(decoded)
    decoded = UpSampling2D((2, 2),dtype=tf.float64)(decoded)

    decoded = Conv2D(ft_bank_baseline, kernel_size=(3, 3), padding='same',dtype=tf.float64)(decoded)
    decoded = Activation('relu',dtype=tf.float64)(decoded)

    decoded = Conv2D(1, kernel_size=(3, 3), padding='same',dtype=tf.float64)(decoded)
    if binary_image:
        outputs = Activation('sigmoid',dtype=tf.float64)(decoded)
    else:
        outputs = decoded

    # instantiate decoder model
    decoder = Model(latent_input, outputs, name='decoder')

    # instantiate VAE model
    pz_mean, pz_log_var = generator(encoder([input_image, input_r])[5])
    outputs = decoder(encoder([input_image, input_r])[2])
    vae = Model([input_image, input_r], [outputs, pz_mean, pz_log_var], name='vae_mlp')

    if binary_image:
        reconstruction_loss = K.mean(binary_crossentropy(tf.cast(input_image, tf.float64), outputs), axis=[1, 2])
    else:
        reconstruction_loss = K.mean(mean_absolute_error(tf.cast(input_image, tf.float64), outputs), axis=[1, 2])

    kl_loss = 1 + z_log_var - pz_log_var - tf.divide(K.square(z_mean - pz_mean), K.exp(pz_log_var)) - tf.divide(K.exp(z_log_var), K.exp(pz_log_var))
    kl_loss = -0.5 * K.sum(kl_loss, axis=-1)
    label_loss = tf.divide(0.5 * tf.square(r_mean - tf.cast(input_r, tf.float64) ), tf.exp(r_log_var)) + 0.5 * r_log_var

    vae_loss = K.mean(reconstruction_loss + kl_loss + label_loss)

    vae.add_loss(vae_loss)
    vae.compile(optimizer='adam')

    return vae, encoder, generator, decoder