import tensorflow as tf
import pandas as pd
import numpy as np
from tensorflow import keras

print(tf.__version__)

EXPLICIT = True

### DATA
training_data = np.random.rand(1000, 90)
train_dataset = tf.data.Dataset.from_tensor_slices((training_data, training_data))
train_dataset = train_dataset.shuffle(1000).batch(100)

### MODEL
x = keras.layers.Input(shape=(90,))
h = keras.layers.Dense(40, activation=tf.nn.relu)(x)
z_mu = keras.layers.Dense(10)(h)
z_sigma = keras.layers.Dense(10, activation=tf.nn.sigmoid)(h)

#########################################################################################
if EXPLICIT:
    eps = tf.random.normal(shape=(10,)) # this works

else:
    batch_size = tf.shape(z_mu)[0]
    dimension = tf.shape(z_mu)[1]
    eps = tf.random.normal(shape=(batch_size, dimension)) # this does NOT work, also tried using shape=(dimension,)

#########################################################################################
z = z_mu + eps * z_sigma

h_decoded = keras.layers.Dense(40, activation=tf.nn.relu)(z)
x_decoded = keras.layers.Dense(90)(h_decoded)

model = keras.models.Model(x, x_decoded)

### LOSS
recon_err = tf.reduce_sum(tf.abs(x - x_decoded), axis=1)
kl_div = -.5 * tf.reduce_sum(1 + 2 * tf.math.log(z_sigma) - tf.square(z_mu) - tf.square(z_sigma), axis=1)
total_loss = tf.reduce_mean(recon_err + kl_div)
model.add_loss(total_loss)

### TRAINING
model.compile(optimizer='adam')
print(model.summary())
model.fit(train_dataset, epochs=5)

### SAVE
keras_file = 'vae_test.h5'
keras.models.save_model(model, keras_file)

### CONVERSION
test_dataset = np.random.rand(100, 90).astype(np.float32)

def representative_dataset_gen():
    for i in range(100):
        yield [test_dataset[i:i+1]]

converter = tf.compat.v1.lite.TFLiteConverter.from_keras_model_file(keras_file)
converter.representative_dataset = representative_dataset_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_file_name = 'vae.tflite'
tflite_model = converter.convert()
open(tflite_file_name, 'wb').write(tflite_model)
