# Created by Hansi at 30/08/2023
import nltk
nltk.download('punkt')
nltk.download('averaged_perceptron_tagger')

import streamlit as st
from accord_nlp.information_extraction.convertor import entity_pairing, graph_building
from accord_nlp.information_extraction.ie_pipeline import InformationExtractor


@st.cache_resource
def init():
    return InformationExtractor()


st.set_page_config(
    page_title='ACCORD NLP Demo',
    initial_sidebar_state='expanded',
    layout='wide',
)

with st.spinner(text="Initialising..."):
    ie = init()


def main():
    st.sidebar.title("ACCORD-NLP")
    st.sidebar.markdown("Extract information from text")
    st.sidebar.markdown(
        "[code](https://github.com/Accord-Project/NLP-Framework)"
    )

    st.header("Input a sentence")

    txt = st.text_area('Sentence')

    # st.write(txt)

    # with st.spinner(text="Processing..."):
    # graph = ie.sentence_to_graph(txt)

    if txt:
        # preprocess
        sentence = ie.preprocess(txt)
        st.write(sentence)

        # NER
        with st.spinner(text="Recognising entities..."):
            ner_predictions, ner_raw_outputs = ie.ner_model.predict([sentence])

        st.write(ner_predictions)

        with st.spinner(text="Extracting relations..."):
            # pair entities to predict their relations
            entity_pair_df = entity_pairing(sentence, ner_predictions[0])
            st.write('entity paired')

            # relation extraction
            re_predictions, re_raw_outputs = ie.re_model.predict(entity_pair_df['output'].tolist())
            entity_pair_df['prediction'] = re_predictions
            st.write(re_predictions)

        with st.spinner(text="Building graph..."):
            # build graph
            graph = graph_building(entity_pair_df, view=False)
            # st.success()

        st.header('Entity-Relation Representation')
        st.graphviz_chart(graph)


if __name__ == '__main__':
    main()