import gradio as gr
import spaces
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

# ==========================================================
# Load Model
# ==========================================================

MODEL_ID = "mirajbhandari/nagrita_ner"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForTokenClassification.from_pretrained(MODEL_ID)

ner = pipeline(
    "token-classification",
    model=model,
    tokenizer=tokenizer,
    aggregation_strategy="simple",
    device=0,
)

# ==========================================================
# Prediction Function
# ==========================================================

@spaces.GPU
def predict(text):
    if not text.strip():
        return "⚠️ Please enter some citizenship text."

    prediction = ner(text)

    merged = []

    # Your original merging logic
    for p in prediction:
        if float(p["score"]) < 0.90:
            continue

        if (
            merged
            and p["word"].startswith("##")
            and p["entity_group"] == merged[-1]["entity_group"]
        ):
            merged[-1]["word"] += p["word"][2:]
        else:
            merged.append({
                "entity_group": p["entity_group"],
                "word": p["word"]
            })

    if not merged:
        return "❌ No entities detected."

    output = "📋 EXTRACTED ENTITIES\n"
    output += "=" * 55 + "\n\n"
    output += f"{'Entity':<20}Value\n"
    output += "-" * 55 + "\n"

    for p in merged:
        output += f"{p['entity_group']:<20}{p['word']}\n"

    return output


# ==========================================================
# UI
# ==========================================================

with gr.Blocks(
    theme=gr.themes.Soft(),
    title="Nepali Citizenship NER",
) as demo:

    gr.Markdown(
        """
# 🇳🇵 Nepali Citizenship NER

Extract structured information from Nepali Citizenship OCR text using a fine-tuned DistilBERT model.

### Supported Fields

- 🆔 Citizenship Number
- 👤 Full Name
- 🚻 Gender
- 🎂 Date of Birth
- 📍 District
- 🏘 Municipality / VDC
- 🏠 Ward Number

💡 Paste OCR text or click one of the examples below.
"""
    )

    with gr.Row():

        with gr.Column(scale=3):

            textbox = gr.Textbox(
                label="📄 Citizenship Text",
                placeholder="Paste the OCR text here...",
                lines=12,
            )

            with gr.Row():
                submit = gr.Button(
                    "🔍 Extract Entities",
                    variant="primary",
                )

                clear = gr.ClearButton(
                    components=[textbox],
                    value="🗑️ Clear"
                )

        with gr.Column(scale=2):

            output = gr.Textbox(
                label="📋 Prediction",
                lines=18,
                interactive=False,
            )

    submit.click(
        fn=predict,
        inputs=textbox,
        outputs=output,
    )

    gr.Examples(
        examples=[
            ["""Citizenship Certificate No.: 73452186
Sex: Male
Full Name: RAM BAHADUR THAPA
Date of Birth (AD): 03 Nov 1992
District: Kaski
Municipality: Pokhara
Ward No.: 11"""],
            ["""Citizenship Certificate No.: 91384572
Sex: Female
Full Name: ANITA GURUNG
Date of Birth (AD): 22 Jan 1996
District: Lalitpur
Municipality: Godawari
Ward No.: 8"""],
        ],
        inputs=textbox,
        label="📚 Try an Example"
    )

    gr.Markdown(
        """
---
### 🚀 About

This demo extracts named entities from Nepali citizenship OCR text using a fine-tuned **DistilBERT** model.

**Model:** `mirajbhandari/nagrita_ner`

Built with ❤️ using 🤗 Transformers, Gradio and Hugging Face Spaces.
"""
    )

demo.launch()