import streamlit as st
import spacy
from spacy import displacy

# Load spaCy model
nlp = spacy.load("en_core_web_sm")

def visualize_entities(doc):
    # Create a list of (start, end, label) tuples for the named entities
    entities = [(ent.start_char, ent.end_char, ent.label_) for ent in doc.ents]

    # Highlight named entities in the input text
    html = displacy.render(doc, style="ent", options={"ents": entities})
    st.markdown(html, unsafe_allow_html=True)

def filter_entities(entities, selected_labels):
    if not selected_labels:
        return entities

    filtered_entities = []
    for entity, label in entities:
        if label in selected_labels:
            filtered_entities.append((entity, label))
    return filtered_entities

def main():
    st.title("Named Entity Recognition App")
    st.write("Enter a text and get named entities!")

    user_input = st.text_area("Enter text:", height=200)
    visualize = st.checkbox("Visualize Entities")
    filter_labels = st.multiselect("Filter Entities by Label:", options=["PERSON", "ORG", "GPE", "DATE"])

    if st.button("Analyze"):
        doc = nlp(user_input)
        entities = [(ent.text, ent.label_) for ent in doc.ents]

        if visualize:
            visualize_entities(doc)

        filtered_entities = filter_entities(entities, filter_labels)

        st.write("Named Entities:")
        for entity, label in filtered_entities:
            st.write(f"Text: {entity}, Entity: {label}")

if __name__ == "__main__":
    main()