# Step 6.1: Define different input components
import gradio as gr

# g. define Dropdown data type
input_module1 = gr.inputs.Dropdown(choices=["KNN", "Softmax", "Deep Neural", "SVM", "Decision Tree","Random Forest"], label = "Model Selection")
# b. define image data type
input_module2 = gr.inputs.Dropdown(choices=["image1", "image2", "image3", "image4","image5"], label = "Sample Image")

# b. define image data type
output_module1 = gr.outputs.Textbox(label = "Predicted Class")

# b. define image data type
output_module2 = gr.outputs.Image(label = "Image")

# Step 6.3: Define a new function that accommodates the input modules.
def multi_inputs(input1, input2):
    import numpy as np
    import pickle
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    import tensorflow as tf
    import scipy.ndimage.interpolation
    (X_train_full, y_train_full), (X_test, y_test) = tf.keras.datasets.mnist.load_data()

    X_valid, X_train = X_train_full[:5000], X_train_full[5000:]
    y_valid, y_train = y_train_full[:5000], y_train_full[5000:]

    if input2 == "image1":
      img = X_test[1]
    if input2 == "image2":
      img = X_test[2]
    if input2 == "image3":
      img = X_test[3]
    if input2 == "image4":
      img = X_test[4]
    if input2 == "image5":
      img = X_test[5]

    import matplotlib.pyplot as plt
    import matplotlib.image as mpimg
    #img = mpimg.imread(img)
    #imag = scipy.ndimage.rotate(input2, 0, reshape=False, cval=0)

    #class_names = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    #"Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
    test_shape = img.shape
    test_width = test_shape[0]
    test_height = test_shape[1]

    image =img.reshape(test_width*test_height)

    if input1 == "KNN":
      loaded_model = pickle.load(open('KNN_Model.sav', 'rb'))
      image = img.reshape((1, -1))
      y_test_pred = loaded_model.predict(image)


    if input1 == "Softmax":
      loaded_model = pickle.load(open('Softmax_Model.sav', 'rb'))
      image = img.reshape((1, -1))
      y_test_pred = loaded_model.predict(image)
      

    if input1 == "SVM":
      loaded_model = pickle.load(open('svm_model.sav', 'rb'))
      image = img.reshape((1, -1))
      y_test_pred = loaded_model.predict(image)

    if input1 == "Deep Neural":
      loaded_model = pickle.load(open('dnn_model.sav', 'rb'))
      image = img.reshape((1, -1))
      y_test_pred = loaded_model.predict(image)
      y_test_pred = np.argmax(y_test_pred,axis=1)

    if input1 == "Decision Tree":
      loaded_model = pickle.load(open('Decision_Tree_model.sav', 'rb'))
      image = img.reshape((1, -1))
      y_test_pred = loaded_model.predict(image)
    
    if input1 == "Random Forest":
      loaded_model = pickle.load(open('randomforest.sav', 'rb'))
      image = img.reshape((1, -1))
      y_test_pred = loaded_model.predict(image)
    
    result = y_test_pred[0]
  
    return result, img
    
# Step 6.4: Put all three component together into the gradio's interface function
gr.Interface(fn=multi_inputs,
             inputs=[input_module1, input_module2],
             outputs=[output_module1,output_module2]
            ).launch( debug = True)