'''This script demonstrates how to build a variational autoencoder with Keras.
Reference: "Auto-Encoding Variational Bayes" https://arxiv.org/abs/1312.6114
'''

import itertools
import sys
import json

import numpy as np
import matplotlib.pyplot as plt
from scipy import sparse
import scipy.io

from keras.layers import Input, Dense, Lambda
from keras.models import Model
from keras import backend as K
from keras import objectives
from keras.datasets import mnist

import pandas
import shelve
import pickle





#batch_size = 16
#original_dim = 784
#latent_dim = 2
#intermediate_dim = 128
#epsilon_std = 0.01
#nb_epoch = 40




def train_vae(x_train,x_dev,x_test,y_train=None,y_dev=None,y_test=None,hidden_size=80,latent_dim=12,batch_size=8,nb_epochs=10,sgd="rmsprop",input_activation = "relu",output_activation = "sigmoid",epsilon_std=0.01):



    def sampling(args):
        z_mean, z_log_std = args
        epsilon = K.random_normal(shape=(batch_size, latent_dim),
                                  mean=0., std=epsilon_std)
        return z_mean + K.exp(z_log_std) * epsilon

    def vae_loss(x, x_decoded_mean):
        xent_loss = objectives.binary_crossentropy(x, x_decoded_mean)
        kl_loss = - 0.5 * K.mean(1 + z_log_std - K.square(z_mean) - K.exp(z_log_std), axis=-1)
        return xent_loss + kl_loss

    original_dim = x_train.shape[1]


    x = Input(batch_shape=(batch_size, original_dim))
    h = Dense(hidden_size, activation=input_activation)(x)
    z_mean = Dense(latent_dim)(h)
    z_log_std = Dense(latent_dim)(h)


    # note that "output_shape" isn't necessary with the TensorFlow backend
    # so you could write `Lambda(sampling)([z_mean, z_log_std])`
    z = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_std])

    # we instantiate these layers separately so as to reuse them later
    decoder_h = Dense(hidden_size, activation=input_activation)
    decoder_mean = Dense(original_dim, activation=output_activation)
    h_decoded = decoder_h(z)
    x_decoded_mean = decoder_mean(h_decoded)


    vae = Model(x, x_decoded_mean)
    vae.compile(optimizer=sgd, loss=vae_loss)

    # train the VAE on MNIST digits
    if y_train is None or y_dev is None or y_test is None :
        y_train = x_train
        y_dev = x_dev
        y_test = x_test

    vae.fit(x_train, y_train,
            shuffle=True,
            nb_epoch=nb_epochs,
            batch_size=batch_size,
            validation_data=(x_dev, y_dev))

    # build a model to project inputs on the latent space
    encoder = Model(x, z_mean)
    pred_train = encoder.predict(x_train, batch_size=batch_size)
    pred_dev = encoder.predict(x_dev, batch_size=batch_size)
    pred_test = encoder.predict(x_test,batch_size=batch_size)
    return [ [ pred_train, pred_dev, pred_test ] ]
# display a 2D plot of the digit classes in the latent space
    #x_test_encoded = encoder.predict(x_test, batch_size=batch_size)
    # build a digit generator that can sample from the learned distribution
    #decoder_input = Input(shape=(latent_dim,))
    #_h_decoded = decoder_h(decoder_input)
    #_x_decoded_mean = decoder_mean(_h_decoded)
    #generator = Model(decoder_input, _x_decoded_mean)
    #x_decoded = generator.predict(z_sample)