#Web application (sample codes from student)
import matplotlib as mpl
import matplotlib.pyplot as plt
import gradio as gr
import numpy as np
import pickle
import tensorflow as tf
from tensorflow import keras
from sklearn.neighbors import KNeighborsClassifier

#Step 1
fashion_mnist = keras.datasets.fashion_mnist
(X_train_full, y_train_full), (X_test, y_test) = fashion_mnist.load_data()

class_names = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
               "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]


# recommended way to save data into image
from PIL import Image
for i in range(6):
  idx = np.random.randint(0,len(X_train_full))
  some_image = X_train_full[idx] # select one image sample
  some_image_label = class_names[y_train_full[idx]] # select one image sample
  some_image = some_image.reshape(28, 28) # reshape from rank-1 tensor (784,) to rank-2 tensor (28,28) 
  im = Image.fromarray(some_image)
  im.save('image'+str(i)+'.jpg')

input_module1 = gr.inputs.Image(label = "test_image", image_mode='L', shape = (28,28))

input_module2 = gr.inputs.Dropdown(choices=["Random Forest", "Decision Tree", "AdaBoost", "Gradient Tree Boosting", "KNN"], label = "Select Algorithm")

output_module1 = gr.outputs.Textbox(label = "Predicted Class")

output_module2 = gr.outputs.Label(label = "Predict Probability")


def fashion_images(input1, input2):
    image = input1.reshape(1, 28*28)/255.0

    import pickle
    with open('knn_model_best.pkl', 'rb') as file:
        best_knn_model = pickle.load(file)

    with open('gradientboost_model_best.pkl', 'rb') as file:
        best_gbt_model = pickle.load(file)

    with open('adaboost_model_best.pkl', 'rb') as file:
        best_adaboost_model = pickle.load(file)

    with open('decision_tree_model_best.pkl', 'rb') as file:
        best_tree_model = pickle.load(file)
        
    with open('random_forest_model_best.pkl', 'rb') as file:
        best_RF_model = pickle.load(file)


    if input2 == 'Random Forest':
        y_test_predicted_proba = best_RF_model.predict_proba(image)
        y_test_predicted_label = best_RF_model.predict(image)
        output = class_names[y_test_predicted_label[0]]

    elif input2 == 'Gradient Tree Boosting':
        y_test_predicted_proba = best_gbt_model.predict_proba(image)
        y_test_predicted_label = best_gbt_model.predict(image)
        output = class_names[y_test_predicted_label[0]]
        
    elif input2 == 'AdaBoost':
        y_test_predicted_proba = best_adaboost_model.predict_proba(image)
        y_test_predicted_label = best_adaboost_model.predict(image)
        output = class_names[y_test_predicted_label[0]]
        
    elif input2 == 'Decision tree':
        y_test_predicted_proba = best_tree_model.predict_proba(image)
        y_test_predicted_label = best_tree_model.predict(image)
        output = class_names[y_test_predicted_label[0]]
        
    elif input2 == 'Random Forest':
        y_test_predicted_proba = best_RF_model.predict_proba(image)
        y_test_predicted_label = best_RF_model.predict(image)
        output = class_names[y_test_predicted_label[0]]

    elif input2 == 'KNN': 
        y_test_predicted_proba = best_knn_model.predict_proba(image)
        y_test_predicted_label = best_knn_model.predict(image)
        output = class_names[y_test_predicted_label[0]]
    
    output_prob = {}
    for name, prob in zip(class_names, y_test_predicted_proba[0]):
        output_prob[name] = prob

    return output, output_prob

# Step 6.4: Put all three component together into the gradio's interface function 
interface = gr.Interface(fn=fashion_images, 
             inputs=[input_module1, input_module2], 
             outputs=[output_module1,output_module2],
             examples=[["image0.jpg", "Random Forest"], 
                       ["image1.jpg", "Decision tree"], 
                       ["image2.jpg", "KNN"]])
interface.launch()