import gradio as gr
from transformers import pipeline

# --------------------------------------
# MODEL CONFIGURATION
# --------------------------------------

model_name = "ELHACHYMI/bert-ner"

default_example = "Bill Gates founded Microsoft in the United States."

example_list = [
    "Elon Musk is the CEO of SpaceX and Tesla.",
    "I am ELHACHYMI Ahmed Yassine and i am from Morocco",
    "Paris is the capital of France.",
    "Othmane AJI student at ENIAD",
]

# Correct mapping from LABEL_X to real entity names
ENTITY_MAP = {
    "LABEL_0": "O",
    "LABEL_1": "PER",
    "LABEL_2": "PER",
    "LABEL_3": "ORG",
    "LABEL_4": "ORG",
    "LABEL_5": "LOC",
    "LABEL_6": "LOC",
    "LABEL_7": "MISC",
    "LABEL_8": "MISC",
}

# Model info for README block
model_info = {
    model_name: {
        "link": f"https://huggingface.co/{model_name}",
        "usage": f"""from transformers import pipeline

ner = pipeline("ner", model="{model_name}", aggregation_strategy="simple")
text = "{default_example}"
print(ner(text))"""
    }
}

# Load NER pipeline
ner_pipeline = pipeline("ner", model=model_name, aggregation_strategy="simple")


# --------------------------------------
# DISPLAY MODEL INFO
# --------------------------------------

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


# --------------------------------------
# HIGHLIGHTED MODE FUNCTION
# --------------------------------------

def analyze_highlight(text):
    results = ner_pipeline(text)
    highlighted = []
    last_idx = 0

    for ent in results:
        start, end = ent["start"], ent["end"]

        # Map LABEL_X → real label
        raw_label = ent["entity_group"]
        label = ENTITY_MAP.get(raw_label, raw_label)

        # Do NOT highlight label "O"
        if label == "O":
            label = None

        # Add normal text before the entity
        if start > last_idx:
            highlighted.append((text[last_idx:start], None))

        # Add entity text (highlight only if label != None)
        highlighted.append((text[start:end], label))

        last_idx = end

    # Add remaining text
    if last_idx < len(text):
        highlighted.append((text[last_idx:], None))

    return highlighted


# --------------------------------------
# GRADIO UI
# --------------------------------------

with gr.Blocks() as demo:
    gr.Markdown("# Named Entity Recognition Demo using BERT")
    gr.Markdown("### Model: ELHACHYMI/bert-ner")

    text_input = gr.Textbox(
        label="Enter Text",
        lines=5,
        value=default_example,
        placeholder="Type a sentence to analyze."
    )

    # Example buttons
    with gr.Row():
        for ex in example_list:
            gr.Button(ex).click(lambda x=ex: x, outputs=text_input)

    # Run button
    analyze_button = gr.Button("Run NER Analysis")

    # Only one output: HighlightedText
    output_highlight = gr.HighlightedText(
        label="NER Result",
        combine_adjacent=True
    )

    # Code + link
    code_output = gr.Code(label="How to use this model")
    link_output = gr.Markdown()

    # Load model info on startup
    demo.load(
        fn=display_model_info,
        inputs=[gr.Textbox(value=model_name, visible=False)],
        outputs=[code_output, link_output]
    )

    # Run analysis
    analyze_button.click(
        analyze_highlight,
        inputs=text_input,
        outputs=output_highlight
    )

demo.launch(ssr_mode=False)