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

#Step 1
fashion_mnist = tk.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"], label = "Select Algorithm")

input_module2 = gr.inputs.Dropdown(choices=['Softmax Regression', 'Neural Network (sklearn)', 'Neural Network (keras) two-layer','Neural Network (keras) three-layer'], 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: # KNN model is too large, ignore here
    #    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':
    '''
    #if input2 == 'Softmax Regression':
    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 
gr.Interface(fn=fashion_images, 
             inputs=[input_module1, input_module2], 
             outputs=[output_module1,output_module2],
             title = "CSCI4750/5750: Build classification pipeline for FashionMNIST",
             examples=[["image0.jpg", "Random Forest"], 
                       ["image1.jpg", "Decision tree"], 
                       ["image2.jpg", "Random Forest"], 
                       ["image3.jpg", "Gradient Tree Boosting"], 
                       ["image4.jpg", "AdaBoost"], 
                       ["image5.jpg", "Random Forest"]]
            ).launch(debug = True)