import keras
from keras import layers
from keras import backend as K
from keras.models import Model
from keras.datasets import mnist

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

# Global Variables
img_shape = (28, 28, 1)
batch_size = 16
latent_dim = 2


def image_input():
    input_img = keras.Input(shape=img_shape)

    x = layers.Conv2D(32, 3, padding="same", activation="relu")(input_img)
    x = layers.Conv2D(64, 3, padding="same", activation="relu", strides=(2, 2))(x)
    x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
    x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
    shape_before_flattening = K.int_shape(x)

    x = layers.Flatten()(x)
    x = layers.Dense(32, activation="relu")(x)

    z_mean = layers.Dense(latent_dim)(x)
    z_log_var = layers.Dense(latent_dim)(x)
    shape_after_flattening = K.int_shape(z_log_var)
    return z_mean, z_log_var, input_img, shape_before_flattening, shape_after_flattening


def sampling(args, latent_dim=2):
    mean, log_var = args
    epsilon = K.random_normal(
        shape=(K.shape(mean)[0], latent_dim), mean=0.0, stddev=1.0
    )
    return mean + K.exp(log_var) * epsilon


def decoder_input(z, shape_before_flattening):

    # This is the input where we will feed `z`.

    decoder_input = layers.Input(shape=K.int_shape(z)[1:])

    # Upsample to the correct number of units
    x = layers.Dense(np.prod(shape_before_flattening[1:]), activation="relu")(
        decoder_input
    )

    # Reshape into an image of the same shape as before our last `Flatten` layer
    x = layers.Reshape(shape_before_flattening[1:])(x)

    # We then apply then reverse operation to the initial
    # stack of convolution layers: a `Conv2DTranspose` layers
    # with corresponding parameters.
    x = layers.Conv2DTranspose(
        32, 3, padding="same", activation="relu", strides=(2, 2)
    )(x)
    x = layers.Conv2D(1, 3, padding="same", activation="sigmoid")(x)
    # We end up with a feature map of the same size as the original input.

    # This is our decoder model.
    decoder = Model(decoder_input, x)

    # We then apply it to `z` to recover the decoded `z`.
    z_decoded = decoder(z)

    return z_decoded, decoder


class CustomVariationalLayer(keras.layers.Layer):
    def vae_loss(self, x, z_decoded):
        x = K.flatten(x)
        z_decoded = K.flatten(z_decoded)
        xent_loss = keras.metrics.binary_crossentropy(x, z_decoded)
        kl_loss = -5e-4 * K.mean(
            1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1
        )
        return K.mean(xent_loss + kl_loss)

    def call(self, inputs):
        x = inputs[0]
        z_decoded = inputs[1]
        loss = self.vae_loss(x, z_decoded)
        self.add_loss(loss, inputs=inputs)
        # We don't use this output.
        return x


def autoencoder_train(input_img, y):
    vae = Model(input_img, y)
    vae.compile(optimizer="rmsprop", loss=None)
    vae.summary()

    # Train the VAE on MNIST digits
    (x_train, _), (x_test, y_test) = mnist.load_data()

    x_train = x_train.astype("float32") / 255.0
    x_train = x_train.reshape(x_train.shape + (1,))
    x_test = x_test.astype("float32") / 255.0
    x_test = x_test.reshape(x_test.shape + (1,))

    vae.fit(
        x=x_train,
        y=None,
        shuffle=True,
        epochs=10,
        batch_size=batch_size,
        validation_data=(x_test, None),
    )


def verify_plots(decoder):
    # Display a 2D manifold of the digits
    n = 15  # figure with 15x15 digits
    digit_size = 28
    figure = np.zeros((digit_size * n, digit_size * n))
    # Linearly spaced coordinates on the unit square were transformed
    # through the inverse CDF (ppf) of the Gaussian
    # to produce values of the latent variables z,
    # since the prior of the latent space is Gaussian
    grid_x = norm.ppf(np.linspace(0.05, 0.95, n))
    grid_y = norm.ppf(np.linspace(0.05, 0.95, n))

    for i, yi in enumerate(grid_x):
        for j, xi in enumerate(grid_y):
            z_sample = np.array([[xi, yi]])
            z_sample = np.tile(z_sample, batch_size).reshape(batch_size, 2)
            x_decoded = decoder.predict(z_sample, batch_size=batch_size)
            digit = x_decoded[0].reshape(digit_size, digit_size)
            figure[
                i * digit_size : (i + 1) * digit_size,
                j * digit_size : (j + 1) * digit_size,
            ] = digit

    plt.figure(figsize=(10, 10))
    plt.imshow(figure, cmap="Greys_r")
    plt.show()


if __name__ == "__main__":
    (
        z_mean,
        z_log_var,
        input_img,
        shape_before_flattening,
        shape_after_flattening,
    ) = image_input()
    z = layers.Lambda(sampling, output_shape=shape_after_flattening)(
        [z_mean, z_log_var]
    )
    z_decoded, decoder = decoder_input(z, shape_before_flattening)
    y = CustomVariationalLayer()([input_img, z_decoded])
    autoencoder_train(input_img, y)
    verify_plots(decoder=decoder)
