import streamlit as st
from models.rubert_MODEL import classify_text
from models.bag_of_words_MODEL import predict
from models.lstm_MODEL import predict_review
import time

class_prefix =  'This review is likely...'


st.title("Movie Review Classification")
st.write("This page will compare three models: Bag of Words/TF-IDF, LSTM, and BERT.")

# Example placeholder for user input
user_input = st.text_area("")


if st.button('Classify with All Models'):
    # Measure and display Bag of Words/TF-IDF prediction time
    start_time = time.time()
    bow_tfidf_result = predict(user_input)
    end_time = time.time()
    st.write(f'{class_prefix} {bow_tfidf_result} according to Bag of Words/TF-IDF. Time taken: {end_time - start_time:.2f} seconds.')
    
    # Measure and display LSTM prediction time
    start_time = time.time()
    lstm_result = predict_review(user_input)
    end_time = time.time()
    st.write(f'{class_prefix} {lstm_result} according to LSTM. Time taken: {end_time - start_time:.2f} seconds.')
    
    # Measure and display ruBERT prediction time
    start_time = time.time()
    rubert_result = classify_text(user_input)
    end_time = time.time()
    st.write(f'{class_prefix} {rubert_result} according to ruBERT. Time taken: {end_time - start_time:.2f} seconds.')


    # Placeholder buttons for model selection
    # if st.button('Classify with BoW/TF-IDF'):
    #     st.write(f'{class_prefix}{predict(user_input)}')
    # if st.button('Classify with LSTM'):
    #     st.write(f'{class_prefix}{predict_review(user_input)}')
    # if st.button('Classify with ruBERT'):
    #     st.write(f'{class_prefix}{classify_text(user_input)}')