import gradio as gr
import joblib

# Load models
models = {
    "Logistic Regression": joblib.load("models/best_model.joblib"),
    "Random Forest": joblib.load("models/random_forest_model.joblib"),
    "KNN": joblib.load("models/trained_knn_model.joblib"),
}

# Load vectorizer
vectorizer = joblib.load("models/vectorizer.joblib")

# Define prediction function
def predict_sentiment(review, model_name):
    # Transform the review text using the vectorizer
    processed_review = vectorizer.transform([review])

    # Select the model
    model = models[model_name]

    # Make predictions
    predicted_class = model.predict(processed_review)[0]
    probabilities = model.predict_proba(processed_review)[0]

    # Define sentiment labels
    sentiment_labels = ["Negative Comment", "Positive Comment"]
    predicted_label = sentiment_labels[predicted_class]

    # Return probabilities as percentages
    positive_percentage = probabilities[1] * 100
    negative_percentage = probabilities[0] * 100

    return predicted_label, positive_percentage, negative_percentage

# Build Gradio interface
with gr.Blocks() as interface:
    gr.Markdown("<h1>Text Classification Models</h1>")
    gr.Markdown("Choose a model and provide a review to see the sentiment analysis results with probabilities displayed as scales.")

    with gr.Row():
        with gr.Column():
            review_input = gr.Textbox(label="Review Comment", placeholder="Type your comment here...")
            model_selector = gr.Dropdown(
                choices=list(models.keys()), label="Select Model", value="Logistic Regression"
            )
            submit_button = gr.Button("Submit")

        with gr.Column():
            sentiment_output = gr.Textbox(label="Predicted Sentiment Class", interactive=False)
            positive_progress = gr.Slider(label="Positive Comment Percentage", minimum=0, maximum=100, interactive=False)
            negative_progress = gr.Slider(label="Negative Comment Percentage", minimum=0, maximum=100, interactive=False)

    submit_button.click(
        predict_sentiment,
        inputs=[review_input, model_selector],
        outputs=[sentiment_output, positive_progress, negative_progress],
    )

# Launch the app
if __name__ == "__main__":
    interface.launch()