import gradio as gr
import tensorflow as tf
from sklearn.preprocessing import LabelEncoder
import spacy
import csv
from timeit import default_timer as timer
from helper_functions import preprocess_single_sentence

with open('examples.txt', 'r') as f:
    example_list = [example.strip() for example in f.readlines()]

with open('label_names.txt', 'r') as f:
    labels = [label.strip() for label in f.readlines()]

encoder = LabelEncoder()
encoder.fit(labels)

model = tf.keras.models.load_model('models/gru_model.keras')

nlp = spacy.load("en_core_web_sm")

def predict(text, flagged=False):
    """
    Make predictions on the given text using the trained model.
    Args:
        text (str): The text to make predictions on.
        flagged (bool): Indicates whether the input was flagged.
    Returns:
        list: A list of predictions.
    """
    start_time = timer()
    processed_text = preprocess_single_sentence(text)
    text_tensor = tf.expand_dims(processed_text, 0)

    probability = model.predict(text_tensor)
    pred_label_with_prob = {labels[i]: float(probability[0][i]) for i in range(len(labels))}
    pred_time = round(timer() - start_time, 5)
    field_names = ['Text', 'Label']

    with open('new_text.txt', 'a') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=field_names)

        # Write header if file is empty
        if csvfile.tell() == 0:
            writer.writeheader()
        writer.writerow({'Text': text, 'Label': encoder.inverse_transform([probability.argmax()])[0]})

    return pred_label_with_prob, pred_time

# Gradio App
inputs = gr.Textbox(lines=5, label="Enter text",
                    placeholder="i like to have the same breathless feeling as a reader eager to see what will happen next")
outputs = [
    gr.Label(num_top_classes=len(labels), label="Predictions"),
    gr.Number(label="Prediction time (s)"),
]
title = 'Sentiment Analysis 🤣😱😡😢'
description = "The sentiment analysis model is a deep learning-based natural language processing (NLP) model designed to analyze and classify the sentiment expressed in text data. It is trained to understand the emotional tone of text and categorize it into predefined sentiment categories such as <b>anger, fear, sadness, and joy.<b><br><br>\
<b>WARNING:</b> Your input will be saved for improving model accuracy. If you believe the model's prediction is incorrect, please click the 'Flag' button.<br>\
<b>Tip:</b> If you assume the output is incorrect, click the 'Flag' button to notify us."
demo = gr.Interface(fn=predict,
                    inputs=inputs,
                    outputs=outputs,
                    title=title,
                    description=description,
                    examples=example_list,
                    allow_flagging=True) 
demo.launch()