# external libraries
import streamlit as st
from transformers import pipeline
    
import pandas as pd

# internal libraries
from config import config
import pipeline


def main():

    st.set_page_config(
        layout="centered",  # Can be "centered" or "wide". In the future also "dashboard", etc.
        initial_sidebar_state="auto",  # Can be "auto", "expanded", "collapsed"
        page_title=config.main_title,  # String or None. Strings get appended with "• Streamlit". 
        page_icon=config.logo_path,  # String, anything supported by st.image, or None.
    )

    if "output" not in st.session_state:
        st.session_state['data'] = pd.read_csv(config.sample_texts_path)
        st.session_state['sample_text'] = None
        generate_text()
        st.session_state["output"] = False
        st.session_state["output_text"] = ""
        st.session_state['inputs'] = {}

    col1, col2, col3 = st.columns(3)
    col1.write(' ')
    col2.image(config.logo_path)
    col3.write(' ')

    st.markdown(f"<h1 style='text-align: center;'>{config.main_title}</h1>", unsafe_allow_html=True)
    st.markdown(f"<h3 style='text-align: center;'>{config.lecture_title}</h3>", unsafe_allow_html=True)

    # topic modelling radio bar
    input_topic_modelling = st.radio(
                        config.topic_modelling_title,
                        config.topic_modelling_answers,
                        horizontal=True)
    st.session_state['inputs']['input_topic_modelling'] = input_topic_modelling

    # input text area
    input_text = st.text_area(config.input_text, st.session_state['sample_text'], height=300)
    st.session_state['inputs']['input_text'] = input_text

    # generate sample text button
    st.button(config.button_text, on_click=generate_text)
    
    # choosing segmenter radio bar
    input_segmenter = st.radio(
                        config.segmenter_title,
                        config.segmenter_answers,
                        horizontal=True)
    st.session_state['inputs']['input_segmenter'] = input_segmenter
    
    # choosing summarizer algorithm radio bar
    input_summarizer = st.radio(
                        config.summarizer_title,
                        config.summarizer_answers,
                        horizontal=True)
    st.session_state['inputs']['input_summarizer'] = input_summarizer
    
    # generating summary button
    col1, col2, col3 = st.columns(3)
    col1.header(' ')
    col2.button(config.generate_text, on_click=generate_summary)
    col3.header(' ')

    if st.session_state["output"]:
        
        TOPICS = [key for key, value in st.session_state["output_text"].items() if key != '#']

        if config.filter_threshold_summaries:
            TOPICS = [key for key in TOPICS if st.session_state["output_text"][key]['summary'] != config.threshold_error]

        st.write(config.output_title)
        options = {}
        for topic in TOPICS:
            option = st.checkbox(topic)
            options[topic] = option

        if len(options) == 0:
            st.warning(config.warning_len_input_text, icon="⚠️")

        for topic, option in options.items():
            if option == True:
                st.text_area(topic, 
                            st.session_state["output_text"][topic]['summary'],
                            disabled=True)

def generate_text():
    df = st.session_state['data']
    df = df[~df['data'].isnull()]
    df = df[df['data'].str.len().gt(100)]
    st.session_state['sample_text'] = df.sample(1)['data'].values[0]

def generate_summary():
    st.session_state["output"] = True

    MODELS = {
        'summarizer':st.session_state['inputs']['input_summarizer'],
        'topic_modelling':st.session_state['inputs']['input_topic_modelling'],
        'segmentizer':st.session_state['inputs']['input_segmenter']
    }

    with st.spinner('Generating the output of Topic Modeling for Summarization...'):
            OUTPUT = pipeline.run(st.session_state['inputs']['input_text'], MODELS)

    st.session_state["output_text"] = OUTPUT
    st.success('Done!')


if __name__ == "__main__":
    main()