import sentencepiece
import streamlit as st
import pandas as pd
import spacy
import roner

example_list = [
    "Ana merge în București.",
    "Președintele Statelor Unite, Joe Biden, a spus, vineri, că va trimite un număr de militari americani în Europa de Est „în curând”, ca urmare a situației tot mai tensionate din Ucraina. Președintele american a spus că „nu va trimite foarte mult” și a exclus din nou posibilitatea desfășurării de trupe în Ucraina, care nu face parte din Alianța Nord-Atlantică.",
    "Deficitul bugetar a crescut semnificativ în ultima lună a anului trecut, cu 24,02 miliarde de lei, după ce, în noiembrie, Finanţele raportau un sold negativ de 4,7% din PIB, respectiv 55,98 miliarde de lei. Pe de altă parte, faţă de 2020, deficitul bugetar a consemnat o scădere importantă, de la 9,61% din PIB, respectiv 101,8 miliarde de lei."
]

st.set_page_config(layout="wide")

st.title("Demo for Romanian NER")

model_list = ['dumitrescustefan/bert-base-romanian-ner']

st.sidebar.header("Select NER Model")
model_checkpoint = st.sidebar.radio("", model_list)

st.sidebar.write("This demo is based on RoNER: 'https://github.com/dumitrescustefan/roner'")
st.sidebar.write("")


st.sidebar.header("Select type of PERSON detection")
named_persons_only_radio = st.sidebar.radio("", ('Proper nouns only', 'All nouns'))

st.sidebar.write("Types of entities detected: 'PERSON', 'ORG', 'GPE', 'LOC', 'NAT_REL_POL', 'EVENT', 'LANGUAGE', 'WORK_OF_ART', 'DATETIME', 'PERIOD', 'MONEY', 'QUANTITY', 'NUMERIC', 'ORDINAL', 'FACILITY'")

st.subheader("Select Text Input Method")
input_method = st.radio("", ('Select from Examples', 'Write or Paste New Text'))
if input_method == 'Select from Examples':
    selected_text = st.selectbox('Select Text from List', example_list, index=0, key=1)
    st.subheader("Text to Run")
    input_text = st.text_area("Selected Text", selected_text, height=128, max_chars=None, key=2)
elif input_method == "Write or Paste New Text":
    st.subheader("Text to Run")
    input_text = st.text_area('Write or Paste Text Below', value="", height=128, max_chars=None, key=2)

@st.cache(allow_output_mutation=True)
def setModel(named_persons_only):
    ner = roner.NER(named_persons_only=named_persons_only)
    return ner

@st.cache(allow_output_mutation=True)
def get_html(html: str):
    WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>"""
    html = html.replace("\n", " ")
    return WRAPPER.format(html)
        
Run_Button = st.button("Run", key=None)
if Run_Button == True:
    
    ner = setModel(named_persons_only = True if named_persons_only_radio.startswith("Proper") else False)
    output = ner(input_text)[0] # only one sentence

    # tabular form
    data = []
    for word in output["words"]:
        if word["tag"]!="O":
            data.append({
                "word": word["text"],
                "entity": word["tag"],
                "start_char": word["start_char"],
                "end_char": word["end_char"],
                "multi_word_entity (cont.)": word["multi_word_entity"]
            })
    df = pd.DataFrame.from_dict(data)
    st.subheader("Recognized Entities")
    st.dataframe(df)


    st.subheader("Spacy Style Display")
    spacy_display = {}
    spacy_display["ents"] = []
    spacy_display["text"] = output["text"]
    spacy_display["title"] = None

    i = 0
    words = output["words"]
    while i < len(words):
        if words[i]["tag"]!="O":
            start = words[i]["start_char"]
            end = words[i]["end_char"]
            label = words[i]["tag"]

            for j in range(i+1,len(words)):
                if words[j]["tag"] == words[i]["tag"] and words[j]["multi_word_entity"] is True:
                    end = words[j]["end_char"]
                    i = j
                j -= 1



            spacy_display["ents"].append({"start": start, "end": end, "label": label})
            print(f"ADD ENTITY: {spacy_display['ents'][-1]}")
        i += 1


    
    entity_list = ['PERSON', 'ORG', 'GPE', 'LOC', 'NAT_REL_POL',
     'EVENT', 'LANGUAGE', 'WORK_OF_ART', 'DATETIME',
     'PERIOD', 'MONEY', 'QUANTITY', 'NUMERIC',
     'ORDINAL', 'FACILITY']

    colors = {
        'PERSON': '#E08989',
        'ORG': '#54478C',
        'GPE': '#2C699A',
        'LOC': '#048BA8',
        'NAT_REL_POL': '#FADAD9',
        'EVENT': '#98654C',
        'LANGUAGE': '#F29E4C',
        'WORK_OF_ART': '#3ADAD9',
        'DATETIME': '#0DB39E',
        'PERIOD': '#83E377',
        'MONEY': '#16DB93',
        'QUANTITY': '#EFEA5A',
        'NUMERIC': '#B9E769',
        'ORDINAL': '#F1C453',
        'FACILITY': '#E0DADA',
    }
    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": entity_list, "colors": colors})
    style = "<style>mark.entity { display: inline-block }</style>"
    st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)