import random

import spacy
import srsly
import streamlit as st

nlp = spacy.load("en_core_web_trf")

# Load pre-processed grants from disk.

grants = list(srsly.read_jsonl("data/processed/entities.jsonl"))

colors = {"GPE": "#5cff84", "LOC": "#5cff84"}
options = {"ents": ["GPE", "LOC"], "colors": colors}

HTML_WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>"""


def render_entities(doc, colors: dict, options: dict) -> str:
    """
    Takes a SpaCy doc
    """

    #if isinstance(doc, spacy.tokens.doc.Doc):
    #    doc = doc.to_json()

    html = spacy.displacy.render(doc, style="ent", options=options)
    html = html.replace("\n", " ")

    return html


st.header("Location Entity Recognition Demo 🔎🌆🌍")

st.subheader("Look for Locations")

if st.button("Show new example", key="text"):
    sample = random.choice(grants)
    doc = nlp(sample["text"])
    html = render_entities(doc, colors, options)
    text = st.text_area("Text input", value=sample["text"], height=200)
    st.write(HTML_WRAPPER.format(html), unsafe_allow_html=True)
else:
    sample = random.choice(grants)
    doc = nlp(sample["text"])
    html = render_entities(doc, colors, options)
    text = st.text_area("Text input", value=sample["text"], height=200, help="Enter text here and click the 'Find Locations' button to search for entities.")
    st.write(HTML_WRAPPER.format(html), unsafe_allow_html=True)