import gradio as gr
from transformers import pipeline


# Model names (keeping it programmatic)
model_names = [
    "c-ho/2026-04-24-crf-classweights-clean",
    "c-ho/2026-04-23-crf-classweights-clean"
]

example_sent = (
    "As a result, Indo-European developed a minimal vowel system combined with a very large consonant inventory including glottalized stops, also grammatical gender and adjectival agreement."
)

# -----------------------
# UI helpers
# -----------------------

# Programmatically build the model info dict
model_info = {
    model_name: {
        "link": f"https://huggingface.co/{model_name}",
        "usage": f"""from transformers import pipeline
ner = pipeline("ner", model="{model_name}", grouped_entities=True)
result = ner("{example_sent}")
print(result)""",
    }
    for model_name in model_names
}

# Load models into a dictionary programmatically for the analyze function
models = {
    model_name: pipeline("ner", model=model_name) #, grouped_entities=True)
    for model_name in model_names
}

def display_model_info(model_name):
    info = model_info[model_name]
    usage_code = info["usage"]
    link = f"[Open model page]({info['link']})"
    return usage_code, link

model_cache = {}
def get_model(model_name):
    if model_name not in model_cache:
        model_cache[model_name] = pipeline(
            "ner",
            model=model_name,
            aggregation_strategy="simple"  # safer for display
        )
    return model_cache[model_name]


# -----------------------
# NER function (SAFE OUTPUT)
# -----------------------

'''
# Function to run NER on input text
def analyze_text(text, model_name):
    ner = models[model_name]
    ner_results = ner(text)
    highlighted_text = []
    last_idx = 0
    for entity in ner_results:
        start = entity["start"]
        end = entity["end"]
        label = entity["entity_group"]
        # Add non-entity text
        if start > last_idx:
            highlighted_text.append((text[last_idx:start], None))
        # Add entity text
        highlighted_text.append((text[start:end], label))
        last_idx = end
    # Add any remaining text after the last entity
    if last_idx < len(text):
        highlighted_text.append((text[last_idx:], None))
    return highlighted_text
'''


def analyze_text(text, model_name):
    ner = get_model(model_name)
    results = ner(text)
    # Convert to safe string output (avoids HighlightedText SSR issues)
    if not results:
        return "No entities found."
    lines = []
    for ent in results:
        word = ent.get("word", "")
        label = ent.get("entity_group", ent.get("entity", "UNK"))
        score = ent.get("score", 0.0)
        lines.append(f"{word} -> {label} ({score:.2f})")
    return "\n".join(lines)

# -----------------------
# UI
# -----------------------
with gr.Blocks() as demo:
    gr.Markdown("# Named Entity Recognition (NER) Demo")
    model_selector = gr.Dropdown(
        choices=model_names,
        value=model_names[0],
        label="Select Model"
    )
    text_input = gr.Textbox(
        label="Input Text",
        lines=5,
        value=example_sent
    )
    btn = gr.Button("Run NER")
    output = gr.Textbox(label="Entities")
    code_output = gr.Code(label="Usage Example")
    link_output = gr.Markdown()
    btn.click(
        analyze_text,
        inputs=[text_input, model_selector],
        outputs=output
    )
    model_selector.change(
        display_model_info,
        inputs=model_selector,
        outputs=[code_output, link_output]
    )
    # Initialize UI on load
    demo.load(
        display_model_info,
        inputs=model_selector,
        outputs=[code_output, link_output]
    )

    
#def greet(name):
#    return "Hello " + name + "!!"

#demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()