# Import dependencies
import tensorflow as tf
from tensorflow import keras

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import gradio as gr
import os

(X_train_full, y_train_full), (X_test, y_test) = keras.datasets.mnist.load_data()

class Sampling(tf.keras.layers.Layer):
    def call(self, inputs):
        mean, log_var = inputs
        return tf.random.normal(tf.shape(log_var)) * tf.exp(log_var / 2) + mean 

custom_objects = {'Sampling': Sampling}

variational_ae = keras.models.load_model('./vae_autoencoder.keras', custom_objects=custom_objects)
variational_encoder = keras.models.load_model('./vae_encoder.keras', custom_objects=custom_objects)
variational_decoder = keras.models.load_model('./vae_decoder.keras', custom_objects=custom_objects)

def get_center_of_mass(df):
    cms = []
    labels = df.label.unique()
    for label in sorted(labels):
        cm = df[df.label == label][df.columns[:-1]].sum() / df[df.label == label].shape[0]
        cms.append(cm.values)
    return cms

_, _, X_test_encoded = variational_encoder(X_test.reshape(-1, 28, 28, 1))
df = pd.concat([pd.DataFrame(X_test_encoded), pd.DataFrame(y_test, columns=['label'])], axis=1)
cms = get_center_of_mass(df)
imgs_cm = [np.clip(variational_decoder(cms[i].reshape(1, -1))[0], 0.0, 255.0) for i in range(10)]

many = []
for i in range(10):#range(4, len(cms) - 5):
    A, B = cms[i - 1].copy(), cms[i].copy()
    versor = (B - A) / np.linalg.norm(B - A)
    d = np.linspace(0, np.linalg.norm(B - A), 100)
    for i in range(100):
        row = A + d[i] * versor
        many.append(row.copy())

many = tf.convert_to_tensor(tf.squeeze(np.array(many)))
images = variational_decoder(many).numpy()
images_converted = []

for x in images:
    new_x = ((x - x.min()) / (x - x.min()).max()) * 255
    new_x = new_x.astype(np.uint8)
    images_converted.append(new_x)

images_converted = np.array(images_converted)


def image_classifier(value):
    return np.clip(((variational_decoder(many[(value * 10)].numpy().reshape(1, -1))[0]) * 255), 0, 255).astype(int)[:, :, 0]

input_value_d = gr.Slider(minimum=2, maximum=99, step=1)
demo = gr.Interface(fn=image_classifier, inputs=input_value_d, outputs="image", live=True)

# Launch the interface
if __name__ == "__main__":
    demo.launch()