import gradio as gr
import pickle
import numpy as np

# Load all the models
with open('random_forest_model.pkl', 'rb') as file:
    random_forest_model = pickle.load(file)

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

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

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

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

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

# Store models in a dictionary for easy access
models = {
    "Random Forest": random_forest_model,
    "Logistic Regression": logistic_model,
    "KNN": knn_yelp_model,
    "SVM Linear": svm_linear,
    "SVM Polynomial": svm_poly,
    "SVM RBF": svm_rbf
}

# Function to predict class probabilities
def predict_class_probabilities(review_text, model_name):
    # Convert review text to the format the model expects (using vectorizer)
    # Assuming you have a vectorizer stored as 'vectorizer.pkl'
    with open('vectorizer.pkl', 'rb') as file:
        vectorizer = pickle.load(file)
    
    # Transform the review text into a vector
    review_vector = vectorizer.transform([review_text])
    
    # Select the model based on user input
    model = models.get(model_name)
    
    if model:
        # Predict class probabilities
        prob = model.predict_proba(review_vector)
        return prob
    else:
        return "Model not found."

# Gradio Interface
interface = gr.Interface(
    fn=predict_class_probabilities, 
    inputs=[
        gr.Textbox(label="Enter your review", placeholder="Type a review here..."),
        gr.Dropdown(
            label="Select a Model", 
            choices=["Random Forest", "Logistic Regression", "KNN", "SVM Linear", "SVM Polynomial", "SVM RBF"]
        )
    ], 
    outputs="json", 
    live=True
)

# Launch the interface
interface.launch()