'''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

  from keras.callbacks import EarlyStopping,Callback

  
  import pandas
  import shelve
  import pickle

  class ZeroStopping(Callback):
      '''Stop training when a monitored quantity has stopped improving.
      # Arguments
          monitor: quantity to be monitored.
          patience: number of epochs with no improvement
              after which training will be stopped.
          verbose: verbosity mode.
          mode: one of {auto, min, max}. In 'min' mode,
              training will stop when the quantity
              monitored has stopped decreasing; in 'max'
              mode it will stop when the quantity
              monitored has stopped increasing.
      '''
      def __init__(self, monitor='val_loss', verbose=0, mode='auto', thresh = 0):
          super(ZeroStopping, self).__init__()
  
          self.monitor = monitor
          self.verbose = verbose
          self.thresh = thresh # is a rythme
  
          if mode not in ['auto', 'min', 'max']:
              warnings.warn('EarlyStopping mode %s is unknown, '
                            'fallback to auto mode.' % (self.mode),
                            RuntimeWarning)
              mode = 'auto'
  
          if mode == 'min':
              self.monitor_op = np.less
          elif mode == 'max':
              self.monitor_op = np.greater
          else:
              if 'acc' in self.monitor:
                  self.monitor_op = np.greater
              else:
                  self.monitor_op = np.less
  
      def on_epoch_end(self, epoch, logs={}):
          current = logs.get(self.monitor)
          if current is None:
              warnings.warn('Zero stopping requires %s available!' %
                            (self.monitor), RuntimeWarning)
  
          if self.monitor_op(current, self.thresh):
              self.best = current
              self.model.stop_training = True

  
  #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,

              verbose = 1,

              batch_size=batch_size,

              validation_data=(x_dev, y_dev)
              #callbacks = [ZeroStopping(monitor='val_loss', thresh=0, verbose=0, mode='min')]

              )

  
      # 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)