import spacy_streamlit
import streamlit as st
import json
from normalizer import process_df
from process_tags import list_ents, no_of_tags, color_creator, scatter_document


def main():
    st.set_page_config(layout='wide')
    st.info(
        "Tags Json is the file we get from exporting tags from d555 tag editor. You may need to format it to be in proper json format. For the demo, use this Tag json file https://github.com/kandysh/odin_json_sheets/blob/main/tag_color.json")
    with st.sidebar:
        uploaded_file = st.file_uploader("Upload the SENTENCE Json", type="json")
        uploaded_color = st.file_uploader("Upload the TAG Json", type="json")
    if uploaded_file and uploaded_color is not None:
        raw_data = json.load(uploaded_file)
        tags_data = json.load(uploaded_color)
        st.title(f'{uploaded_file.name.split(".")[0].upper()}')
        df_list = [process_df(data) for data in raw_data]
        st.plotly_chart(scatter_document(df_list, tags_data), use_container_width=True)
        key = 0
        for df in df_list:
            ents = list_ents(df)
            tags = list(no_of_tags(df).keys())
            doc = [{
                "text": ' '.join(df['words']),
                "ents": ents,
                "title": None
            }]
            st.text(f"Sentence {key}")
            spacy_streamlit.visualize_ner(
                doc,
                labels=tags,
                show_table=False,
                title=None,
                manual=True,
                displacy_options={
                    "colors": color_creator(tags_data["NER"])
                },
                key=f"{key}"
            )
            key += 1


if __name__ == "__main__":
    main()