'''
Sentence model, inspired by the Keras VAE tutorial.

Reference: "Auto-Encoding Variational Bayes" https://arxiv.org/abs/1312.6114
'''
import numpy as np
import tensorflow as tf

from keras.layers import Input, Embedding, Dense, Lambda, LSTM, RepeatVector, TimeDistributed
from keras.models import Model
from keras import backend as K
from keras import objectives
from keras.preprocessing import text, sequence

from utils import load_data, to_one_hot

batch_size = 30
nb_epoch = 50

max_sequence_len = 25 # Required for tensorflow!
lstm_state_dim = 256
latent_dim = 128
intermediate_dim = 256

text_input = Input(shape=(max_sequence_len,), dtype='int32', name='text_input')
embedding = Embedding(1000, 64, input_length=max_sequence_len, mask_zero=True, dropout=0.3)(text_input)
rnn_encoded = LSTM(lstm_state_dim, dropout_W=0.3)(embedding)

h = Dense(intermediate_dim, activation='relu')(rnn_encoded)
z_mean = Dense(latent_dim, name="z_mean_dense")(h)
z_log_var = Dense(latent_dim, name="z_log_var_dense")(h)


def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(batch_size, latent_dim), mean=0.)
    return z_mean + K.exp(z_log_var / 2) * epsilon

z = Lambda(sampling)([z_mean, z_log_var])

h_decoded = Dense(intermediate_dim, activation='relu', name="h_decoded_dense", input_shape=(latent_dim,))(z)
x_decoded_mean = Dense(lstm_state_dim, activation='relu', name="x_decoded_mean_dense")(h_decoded)

rnn_decoded = LSTM(1000, return_sequences=True, dropout_W=0.3)(RepeatVector(max_sequence_len)(x_decoded_mean))
text_output = TimeDistributed(Dense(1000, activation='softmax'))(rnn_decoded)


def vae_loss(text_true_onehot, text_predicted_onehot):
    xent_loss = K.mean(objectives.categorical_crossentropy(text_true_onehot, text_predicted_onehot), axis=-1)
    kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
    return xent_loss + kl_loss

train, dev, scores = load_data()

x_train = [text.one_hot(x, n=1000, lower=True, split=" ") for x in train[0]]
x_train = sequence.pad_sequences(x_train, maxlen=max_sequence_len).astype(int)

#print(x_train[:2])

x_train_one_hot = np.asarray(list(map(to_one_hot, x_train)))

#print(x_train_one_hot[:2])

vae = Model(text_input, text_output)
vae.compile(optimizer='adam', loss=vae_loss, metrics=['accuracy'])

vae.fit(x_train, x_train_one_hot,
        shuffle=True,
        nb_epoch=nb_epoch,
        batch_size=batch_size)

vae.save("model-latest.keras")

#vae.evaluate(x_test, x_test, batch_size=256)

# build a model to project inputs on the latent space
#encoder = Model(text_input, z_mean)
