import streamlit as st
from functions_preprocess import LinguisticPreprocessor, download_if_non_existent, CNN
import pickle
import nltk
import torch
nltk.download('stopwords')
nltk.download('punkt')
download_if_non_existent('corpora/stopwords', 'stopwords')
download_if_non_existent('taggers/averaged_perceptron_tagger', 'averaged_perceptron_tagger')
download_if_non_existent('corpora/wordnet', 'wordnet')


#################################################################### Streamlit interface
st.title("Movie Reviews: An NLP Sentiment analysis")

#################################################################### Cache the model loading

@st.cache_data()
def load_model():
  model_pkl_file = "sentiment_model.pkl"  
  with open(model_pkl_file, 'rb') as file:  
    model = pickle.load(file)
  return model

def load_cnn():
  model = CNN(16236, 300, 128, [3, 8], 0.5, 2)
  model.load_state_dict(torch.load('model_cnn.pkl'))
  model.eval()

  return model

def predict_sentiment(text, model, torch=False):
    if torch == True:
        processed_text = processor.transform(text)
        with torch.no_grad():  # Ensure no gradients are computed
            prediction = model(processed_text)  # Get raw model output
            # Convert output to probabilities
            probs = torch.softmax(prediction, dim=1)
            # Get the predicted class
            pred_class = torch.argmax(probs, dim=1)
        return pred_class.item()  # Return the predicted class as a Python int
    else:
        processor.transform(text)
        prediction = model.predict([text])
        return prediction
        


model_1 = load_model()
model_2 = load_cnn()
processor = LinguisticPreprocessor()


############################################################# Text input

with st.expander("Model 1: SGD Classifier"):
    st.markdown("Give it a go by writing a positive or negative text, and analyze it!")

    # Text input inside the expander
    user_input = st.text_area("Enter text here...", key='model1_input')
    if st.button('Analyze', key='model1_button'):
        # Displaying output
        result = predict_sentiment(user_input, model_1)
        if result >= 0.5:
            st.write('The sentiment is: Positive 😀', key='model1_poswrite')
        else:
            st.write('The sentiment is: Negative 😞', key='model1_negwrite')

with st.expander("Model 2: CNN Sentiment analysis"):
    st.markdown("Give it a go by writing a positive or negative text, and analyze it!")

    # Text input inside the expander
    user_input = st.text_area("Enter text here...", key='model2_input')
    if st.button('Analyze', key='model2_button'):
        # Displaying output
        result = predict_sentiment(user_input, model_2, torch=True)
        if result >= 0.5:
            st.write('The sentiment is: Positive 😀', key='model2_poswrite')
        else:
            st.write('The sentiment is: Negative 😞', key='model2_negwrite')

st.caption("Por @efeperro.")