"""
    @author : Sakshi Tatak
"""

# Imports
import pandas as pd
import streamlit as st

from predict_flair import SentimentClassifier as FlairSentimentClassifier
from predict_ml import predict as predict_ml
from predict_setfit import SentimentClassifier as SetFitSentimentClassifier
from predict_spacy import SentimentClassifier as SpacySentimentClassifier

st.set_page_config(layout = 'wide')
st.title('SetFit, Flair, SpaCy, Naive Bayes Sentiment Classifiers')

if 'flair_model' not in st.session_state:
    st.session_state['flair_model'] = None

if 'spacy_model' not in st.session_state:
    st.session_state['spacy_model'] = None

if 'setfit_model' not in st.session_state:
    st.session_state['setfit_model'] = None

if 'results' not in st.session_state:
    st.session_state['results'] = pd.DataFrame(columns = ['model', 'query', 'sentiment', 'confidence'])

def main():
    model_name = st.selectbox('Select Model', options = ['SetFit', 'Naive Bayes', 'Flair', 'SpaCy'])
    if model_name == 'SetFit':
        if st.session_state.setfit_model is None:
            with st.spinner('Loading SetFit classifier ...'):
                st.session_state.setfit_model = SetFitSentimentClassifier()
            st.success('SetFit classifier loaded successfully!')
        model = st.session_state.setfit_model

    if model_name == 'Flair':
        if st.session_state.flair_model is None:
            with st.spinner('Loading Flair classifier ...'):
                st.session_state.flair_model = FlairSentimentClassifier()
            st.success('Flair classifier loaded successfully!')
        model = st.session_state.flair_model

    if model_name == 'SpaCy':
        if st.session_state.spacy_model is None:
            with st.spinner('Loading SpaCy classifier'):
                st.session_state.spacy_model = SpacySentimentClassifier()
            st.success('Spacy classifier loaded successfully!')
        model = st.session_state.spacy_model

    text = st.text_area('Input text', value = 'This is insane haha!')

    if st.button('Compute sentiment'):
        if model_name != 'Naive Bayes':
            with st.spinner(f'Predicting with {model_name} ...'):
                sentiment, conf = model.predict(text)
        else:
            with st.spinner('Predicting with Naive Bayes ...'):
                sentiment, conf = predict_ml(text)

        st.success(sentiment + ', ' + str(conf))
        df = st.session_state.results
        df.loc[len(df)] = [model_name, text, sentiment, conf]
        st.table(df)



if __name__ == '__main__':
    main()