# 1. Import the required packages
import torch
import gradio as gr

from typing import Dict, Tuple
from transformers import pipeline

# 2. Define function to use our model on given text
def blood_request_classifier(text: str) -> Tuple[str, Dict[str, float]]:
    # Set up text classification pipeline
    classifier = pipeline(
        task="text-classification",
        model="AshenFdo/emergency_blood_request_classifier",
        device="cuda" if torch.cuda.is_available() else "cpu",
        top_k=None
    )

    # Get outputs from pipeline
    outputs = classifier(text)[0]

    # Build probability scores dict + find top prediction
    prob_scores = {}
    top_label = ""
    top_score = 0.0

    for item in outputs:
        label = "🚨 Emergency" if item["label"] == "LABEL_1" else "✅ Non-Emergency"
        prob_scores[label] = round(item["score"], 4)
        if item["score"] > top_score:
            top_score = item["score"]
            top_label = label

    # Build a nice verdict string
    verdict = f"{top_label}  —  Confidence: {round(top_score * 100, 2)}%"

    return verdict, prob_scores

# 3. Create a Gradio interface
description = """
A text classifier to determine whether a blood donation request is an **emergency** or **non-emergency**.

Fine-tuned from [DistilBERT](https://huggingface.co/distilbert/distilbert-base-uncased) on a 
[synthetic blood request urgency dataset](https://huggingface.co/datasets/AshenFdo/synthetic_blood_request_urgency_dataset).

See [source code on GitHub](https://github.com/AshenFdo/Blood-Request-Emergency-Classification-Model).
"""

demo = gr.Interface(
    fn=blood_request_classifier,
    inputs=gr.Textbox(
        lines=4,
        placeholder="Enter a blood request message here...",
        label="Blood Request Text"
    ),
    outputs=[
        gr.Textbox(label="🏷️ Verdict"),          # Shows label + confidence %
        gr.Label(num_top_classes=2, label="📊 Probability Scores")   # Shows both class probs
    ],
    title="🩸 Emergency Blood Request Classifier",
    description=description,
    examples=[
        ["Patient is in critical condition after surgery and urgently needs O- blood immediately or they may not survive."],
        ["Hi, I am looking for a B+ blood donor for my father's scheduled knee replacement surgery next month."],
        ["URGENT: Accident victim in ER needs AB+ blood NOW. Lives at stake, please respond immediately!"],
        ["Our hospital is planning a blood donation camp next Saturday. All blood types welcome."],
        ["A newborn baby in the ICU critically needs O+ blood within the next hour. Please help!"],
    ],
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

# 4. Launch the interface
if __name__ == "__main__":
    demo.launch()