import pickle
import numpy as np
import os
os.system("pip install scikit-learn")
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier
os.system("pip install gradio")
import gradio as gr

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

def fashion_images(image, model):
    print("image: ",image.shape)
    numpy_image = image.reshape(1, 28*28)
    def select_model(model):
      match (model):
        case 'Softmax Regression':
          return pickle.load(open('random_forest_model_best.pkl', 'rb'))
        case 'Neural Network (sklearn) 1':
          return pickle.load(open('random_forest_model_best.pkl', 'rb'))
        case 'Neural Network (sklearn) 2':
          return pickle.load(open('random_forest_model_best.pkl', 'rb'))
        case 'Neural Network (keras) three-layer':
          return pickle.load(open('keras_sequential', 'rb'))
        case _:
          raise BaseException('Model not valid. Please pick a valid model.')
    user_model = select_model(model)
    predicted_class = user_model.predict(numpy_image)[0]
    print("predicted_class: ",predicted_class)
    predicted_class = class_names[predicted_class]
    predicted_proba = user_model.predict_proba(numpy_image)[0]
    print("predicted_proba: ",predicted_proba)
    predicted_proba = {class_names: float(prob) for class_names,prob in zip(class_names, predicted_proba) }
    return predicted_class, predicted_proba


input_module1 = gr.Image(label = "test_image", image_mode='L', shape=(28, 28))
input_module2 = gr.Dropdown(choices=['Softmax Regression', 'Neural Network (sklearn) 1', 'Neural Network (sklearn) 2','Neural Network (keras) three-layer'], label = "Select Algorithm")

output_module1 = gr.Textbox(label = "Predicted Class")
output_module2 = gr.Label(label = "Predict Probability")

gr.Interface(
    fn=fashion_images,
    inputs=[input_module1, input_module2],
    outputs=[output_module1, output_module2]
).launch(debug=True)