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

example_list = [
    "George merge cu trenul Cluj - Timișoara de ora 6:20.",
    "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.",
    "Din anul 2050, doar 10 din cele 21 de locații care au găzduit Jocurile Olimpice de iarnă din 1924 și până în prezent vor avea zăpada naturală.",
    "Al Doilea Război Mondial a fost un război global care a durat între anii 1939 - 1945.",
]

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.header("Select type of PERSON detection")
named_persons_only_radio = st.sidebar.radio("", ('Proper nouns only', 'All nouns'))


st.sidebar.write("This demo is based on RoNER: 'https://github.com/dumitrescustefan/roner'")
st.sidebar.write("")
st.sidebar.write("The NER corpus used is: 'https://github.com/dumitrescustefan/ronec'")
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
                else:
                    break

            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': '#8477AC',
        'GPE': '#4C89AA',
        'LOC': '#34ABF8',
        '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)