import streamlit as st
import spacy
from transformers import pipeline
from PIL import Image

nlp = spacy.load("./models/en_core_web_sm")

st.title("SpaGAN Demo")
st.write("Enter a text, and the system will highlight the geo-entities within it.")

# Define a color map and descriptions for different entity types
COLOR_MAP = {
    'FAC': ('red', 'Facilities (e.g., buildings, airports)'),
    'ORG': ('blue', 'Organizations (e.g., companies, institutions)'),
    'LOC': ('purple', 'Locations (e.g., mountain ranges, water bodies)'),
    'GPE': ('green', 'Geopolitical Entities (e.g., countries, cities)')
}

# Display the color key with descriptions
st.write("**Color Key:**")
for label, (color, description) in COLOR_MAP.items():
    st.markdown(f"- **{label}**: <span style='color:{color}'>{color}</span> - {description}", unsafe_allow_html=True)

# Text input
user_input = st.text_area("Input Text", height=200)

# Process the text when the button is clicked
if st.button("Highlight Geo-Entities"):
    if user_input.strip():
        # Process the text using spaCy
        doc = nlp(user_input)
        
        # Generate highlighted text with different colors for each entity type
        highlighted_text = ""
        last_pos = 0

        for ent in doc.ents:
            color = COLOR_MAP.get(ent.label_, ('black', ''))[0]  # Default to black if label not in map

            # Add text before the entity
            highlighted_text += user_input[last_pos:ent.start_char]

            # Add the highlighted entity text
            highlighted_text += f"<span style='color:{color}; font-weight:bold'>{ent.text}</span>"

            # Update the position
            last_pos = ent.end_char

        # Add any remaining text after the last entity
        highlighted_text += user_input[last_pos:]

        # Display the highlighted text with HTML support
        st.markdown(highlighted_text, unsafe_allow_html=True)
    else:
        st.error("Please enter some text.")