import gradio as gr
from transformers import pipeline
from collections import defaultdict

ner_pipeline = pipeline("ner", model="syubraj/spanish_bert_based_ner")

ENTITY_COLORS = {
    "B-PER": "#FF5733",  
    "I-PER": "#FF8D1A",  
    "B-LOC": "#33FF57",  
    "I-LOC": "#2EB82E",  
    "B-ORG": "#3380FF",  
    "I-ORG": "#1A66FF",  
    "O": "#E0E0E0"       
    }

def highlight_entities(text):
    """
    This function takes a Spanish text input and applies Named Entity Recognition (NER) using a pre-trained model.
    The recognized entities are highlighted with distinct colors based on their categories.
    A legend is also generated to indicate the meaning of each color.
    
    Parameters:
    text (str): The input Spanish text.
    
    Returns:
    str: The HTML formatted output with highlighted entities and a legend.
    """
    entities = ner_pipeline(text)
    formatted_text = ""
    last_idx = 0
    
    for ent in entities:
        word = ent['word']
        start, end = ent['start'], ent['end']
        entity_label = ent['entity']
        color = ENTITY_COLORS.get(entity_label, "#E0E0E0")
        
        formatted_text += text[last_idx:start] + f'<span style="background-color:{color}; padding:2px; border-radius:5px;">{word}</span>'
        last_idx = end
    
    formatted_text += text[last_idx:]
    
    # Generate legend
    legend_html = "<div><b>Legend:</b><br>"
    for label, color in ENTITY_COLORS.items():
        legend_html += f'<span style="background-color:{color}; padding:2px 5px; border-radius:5px; margin-right:5px;">{label}</span>'
    legend_html += "</div><br>"
    
    return legend_html + formatted_text

example_sentences = [
    "Elon Musk vive en Estados Unidos y es dueño de SpaceX, Tesla y Starlink.",
    "Lionel Messi juega para el Inter Miami y ha ganado múltiples Balón de Oro.",
    "Amazon es una de las empresas tecnológicas más grandes con sede en Seattle, EE.UU.",
    "Madrid es la capital de España y alberga el famoso museo del Prado.",
    "Shakira nació en Colombia y es una de las artistas más reconocidas a nivel mundial."
]

def example_selector(example):
    return highlight_entities(example)

grn = gr.Interface(
    fn=highlight_entities,
    inputs=gr.Textbox(label="Enter Spanish Text"),
    outputs=gr.HTML(label="NER Highlighted Text"),
    title="Spanish Named Entity Recognition",
    description="This interactive demo performs Named Entity Recognition (NER) on Spanish text. Recognized entities such as persons, locations, and organizations are highlighted in distinct colors for better readability. A legend is provided to help interpret the color coding.",
    examples=example_sentences,
    allow_flagging="never"
)

grn.launch()