import gradio as gr
import numpy as np
import tensorflow as tf
import pickle
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import StratifiedShuffleSplit

fashion_mnist = tf.keras.datasets.fashion_mnist.load_data()
(X_train_full, y_train_full), (X_test, y_test) = fashion_mnist
X_train, y_train = X_train_full[:-5000], y_train_full[:-5000]
X_valid, y_valid = X_train_full[-5000:], y_train_full[-5000:]

data_split = StratifiedShuffleSplit(n_splits=1, test_size=0.9, random_state=0) # split data one time into two parts with ratio 10%/90%
for train_index, test_index in data_split.split(X_train, y_train):
    small_x_train, small_y_train = X_train[train_index], y_train[train_index]

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

class_labels = ['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')
    
train_size, width, height = small_x_train.shape
small_x_train_reshape = small_x_train.reshape(train_size, width * height)

best_softmax_reg = LogisticRegression(C=0.5, random_state=42)
best_softmax_reg.fit(small_x_train_reshape, small_y_train)

best_nn = MLPClassifier(hidden_layer_sizes=(100, 100), activation='relu', solver='lbfgs')
best_nn.fit(small_x_train_reshape, small_y_train)
'''
input_module1 = gr.Image(label = "test_image", image_mode='L', shape = (28,28))
input_module2 = gr.Dropdown(choices=['Softmax Regression', 'Neural Network (sklearn)'], label = "Select Algorithm")
output_module1 = gr.Textbox(label = "Predicted Class")
output_module2 = gr.Label(label = "Predicted Probability")
'''
input_module1 = gr.Image(label = "test_image", image_mode='L')

# Specify method dropdown menu
input_module2 = gr.Dropdown(choices=['Softmax Regression', 'Neural Network (sklearn)', 'Neural Network (keras) two-layer','Neural Network (keras) three-layer'], label = "Select Algorithm")

# Specify output 1
output_module1 = gr.Textbox(label = "Predicted Class")

# Specify output 2
output_module2 = gr.Label(label = "Predict Probability")



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

best_nn = best_softmax_reg


def fashion_images(input1, input2): # input 1 = image, input 2 = model selection
    numpy_image = input1.reshape(1, 28*28) 
    
    if input2 == 'Softmax Regression':
        probabilities = best_softmax_reg.predict_proba(numpy_image)[0]
    else:
        probabilities = best_nn.predict_proba(numpy_image)[0]
    predicted_class = class_labels[np.argmax(probabilities)]

    print("probabilities: ",probabilities)
    print("class_labels: ",class_labels)
    predicted_probabilities = {class_label: float(prob) for class_label, prob in zip(class_labels, probabilities)}

    return predicted_class, predicted_probabilities


# 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)