import numpy as np
import gradio as gr
from tensorflow.keras.models import load_model
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt

def get_model(path):
    saved_VRAE_model = load_model(path, compile=False)
    
    # return encoder, regressor
    return saved_VRAE_model.layers[1], saved_VRAE_model.layers[2]

def predict(subset, engine):
  # ------------------------ Prediction ------------------------
  path = './models/'+subset+'.h5'
  y_train = np.load('./data/'+subset+'_y_train.npy', allow_pickle=True)
  x_test = np.load('./data/'+subset+'_x_test.npy', allow_pickle=True)
  
  # Get model
  encoder, regressor = get_model(path)
  
  # latent space in train
  fig = plt.figure()
  z_train = np.load('./data/'+subset+'_z_train.npy', allow_pickle=True)
  plt.scatter(z_train[:, 0], z_train[:, 1], c=y_train, cmap='RdYlGn')
  clb = plt.colorbar()
  clb.set_label('RUL', labelpad=-30, y=1.05, rotation=0)
  z, _, _  = encoder.predict(x_test[engine].reshape(1, x_test[engine].shape[0], x_test[engine].shape[1]))
  plt.scatter(z[:, 0], z[:, 1], c='blue', marker='x', s=60)
  plt.xlabel('z - dim 1')
  plt.ylabel('z - dim 2')
  
  pred = regressor.predict(z)[0][0]
  
  return int(pred), fig
  
iface = gr.Interface(
	fn=predict,
	inputs=[gr.inputs.Radio(["FD001", "FD002", "FD003", "FD004"]), gr.inputs.Slider(1, 100, step=1)],
	title="RUL Estimation for CMAPSS dataset",
	description="Enter subset and engine number to predict the remaining useful life (RUL)",
	outputs=[gr.outputs.Label(label="Prediction"), gr.outputs.Plot(type="auto", label="Latent space")],
	allow_screenshot=False,
  theme="dark-peach",
	layout="unaligned")
iface.launch()