import time
import gradio as gr
from gliner import GLiNER

model = GLiNER.from_pretrained("Ihor/gliner-biomed-base-v1.0")

MAX_LABELS = 12

PALETTE = [
    "#FF6B6B", "#4ECDC4", "#45B7D1", "#FFA07A", "#98D8C8", "#F7DC6F",
    "#BB8FCE", "#85C1E9", "#F0B27A", "#76D7C4", "#F1948A", "#82E0AA",
]

DEFAULT_LABELS = ["patient_name", "age", "sex", "symptom", "diagnosis", "medication", "vital_sign", "procedure"]


def filter_choices(selected):
    return gr.Dropdown(choices=[l for l in DEFAULT_LABELS if l not in selected])


def extract(text, labels_list, threshold):
    labels = [l for l in labels_list if l][:MAX_LABELS]
    if not labels or not text.strip():
        return None, [], ""

    color_map = {l: PALETTE[i % len(PALETTE)] for i, l in enumerate(labels)}

    start = time.perf_counter()
    entities = model.predict_entities(text, labels, threshold=threshold)
    latency_ms = (time.perf_counter() - start) * 1000

    entities = sorted(entities, key=lambda e: e["start"])

    hl_entities = [{"entity": e["label"], "start": e["start"], "end": e["end"]} for e in entities]
    table = [[e["label"], e["text"], f"{e['score']:.2f}"] for e in entities]

    return (
        gr.HighlightedText(value={"text": text, "entities": hl_entities}, color_map=color_map),
        table,
        f"{latency_ms:.1f} ms | {len(entities)} entities",
    )


EXAMPLES = [
    [
        """Patient: Jane Doe, 58-year-old female.
Chief Complaint: Chest pain and shortness of breath for 2 days.

History of Present Illness:
Patient presents with substernal chest pain radiating to the left arm,
rated 7/10, worsening with exertion. She reports associated dyspnea and
diaphoresis. She has a history of Type 2 Diabetes Mellitus diagnosed in
2015 and Hypertension diagnosed in 2018.

Current Medications:
- Metformin 1000mg PO BID for diabetes
- Lisinopril 20mg PO daily for hypertension
- Aspirin 81mg PO daily for cardiac prophylaxis

Vitals: BP 158/92, HR 96, SpO2 94%, Temp 98.6F

Assessment:
1. Acute coronary syndrome - rule out myocardial infarction
2. Uncontrolled hypertension
3. Type 2 Diabetes Mellitus - stable on current regimen

Plan:
- Stat ECG and troponin levels
- Start Heparin drip 18 units/kg/hr IV
- Nitroglycerin 0.4mg sublingual PRN chest pain
- Cardiology consult
- Continue home medications""",
        DEFAULT_LABELS,
        0.4,
    ],
    [
        """DISCHARGE SUMMARY
Patient: Robert Chen, 72-year-old male.
Admission Date: 2024-01-15. Discharge Date: 2024-01-19.

Principal Diagnosis: Community-acquired pneumonia, right lower lobe.
Secondary Diagnoses: COPD, Atrial fibrillation.

Hospital Course:
Patient admitted with fever 101.8F, productive cough with purulent sputum,
and oxygen saturation of 88% on room air. Chest X-ray confirmed right lower
lobe consolidation. Started on Ceftriaxone 1g IV daily and Azithromycin
500mg PO daily. Supplemental O2 via nasal cannula at 3L/min.

Discharge Medications:
- Amoxicillin-Clavulanate 875mg PO BID x 5 days
- Albuterol inhaler 2 puffs q4-6h PRN
- Warfarin 5mg PO daily
- Metoprolol 50mg PO BID

Follow-up: Pulmonology clinic in 2 weeks. Repeat chest X-ray in 6 weeks.""",
        DEFAULT_LABELS,
        0.4,
    ],
    [
        """ED Note - 03/10/2024 22:45
Chief Complaint: Laceration to right hand.

HPI: 34-year-old male presents after cutting his right palm on broken glass
approximately 1 hour ago. Reports moderate bleeding controlled with direct
pressure. Denies numbness or weakness in fingers. No foreign body sensation.
Tetanus up to date.

Exam: 3cm linear laceration to right thenar eminence, clean edges, no tendon
involvement, neurovascular intact distally.

Procedure: Wound irrigated with normal saline. Repaired with 4-0 nylon,
5 interrupted sutures. Sterile dressing applied.

Disposition: Home with wound care instructions. Suture removal in 10 days.""",
        DEFAULT_LABELS + ["body_part", "wound"],
        0.4,
    ],
]

with gr.Blocks(title="GLiNER Biomedical NER") as demo:
    gr.Markdown("# GLiNER Biomedical NER\nZero-shot named entity recognition with `gliner-biomed-base-v1.0`")

    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(label="Clinical Text", lines=12)
            labels_input = gr.Dropdown(
                label="Entity Labels",
                choices=DEFAULT_LABELS,
                value=DEFAULT_LABELS,
                multiselect=True,
                allow_custom_value=True,
                max_choices=MAX_LABELS,
            )
            threshold = gr.Slider(0.0, 1.0, value=0.4, step=0.05, label="Confidence Threshold")
            run_btn = gr.Button("Extract", variant="primary")

        with gr.Column(scale=3):
            latency_output = gr.Textbox(label="Latency")
            highlight_output = gr.HighlightedText(label="Entities", combine_adjacent=False, show_legend=True)
            table_output = gr.Dataframe(headers=["Label", "Text", "Score"], label="Extracted Entities")

    labels_input.change(filter_choices, inputs=[labels_input], outputs=[labels_input])
    run_btn.click(extract, inputs=[text_input, labels_input, threshold], outputs=[highlight_output, table_output, latency_output])
    gr.Examples(EXAMPLES, inputs=[text_input, labels_input, threshold])

demo.launch()