import gradio as gr
from transformers import pipeline

# 1. Load the model
pipe = pipeline("text-classification", model="shanover/symps_disease_bert_v3_c41")

# 2. Define the missing label mapping (Alphabetical order + Vertigo first)
id2label = {
    "LABEL_0": "(Vertigo) Paroxysmal Positional Vertigo",
    "LABEL_1": "AIDS",
    "LABEL_2": "Acne",
    "LABEL_3": "Alcoholic hepatitis",
    "LABEL_4": "Allergy",
    "LABEL_5": "Arthritis",
    "LABEL_6": "Bronchial Asthma",
    "LABEL_7": "Cervical spondylosis",
    "LABEL_8": "Chicken pox",
    "LABEL_9": "Chronic cholestasis",
    "LABEL_10": "Common Cold",
    "LABEL_11": "Dengue",
    "LABEL_12": "Diabetes",
    "LABEL_13": "Dimorphic hemmorhoids(piles)",
    "LABEL_14": "Drug Reaction",
    "LABEL_15": "Fungal infection",
    "LABEL_16": "GERD",
    "LABEL_17": "Gastroenteritis",
    "LABEL_18": "Heart attack",
    "LABEL_19": "Hepatitis A",
    "LABEL_20": "Hepatitis B",
    "LABEL_21": "Hepatitis C",
    "LABEL_22": "Hepatitis D",
    "LABEL_23": "Hepatitis E",
    "LABEL_24": "Hypertension",
    "LABEL_25": "Hyperthyroidism",
    "LABEL_26": "Hypoglycemia",
    "LABEL_27": "Hypothyroidism",
    "LABEL_28": "Impetigo",
    "LABEL_29": "Jaundice",
    "LABEL_30": "Malaria",
    "LABEL_31": "Migraine",
    "LABEL_32": "Osteoarthristis",
    "LABEL_33": "Paralysis (brain hemorrhage)",
    "LABEL_34": "Peptic ulcer diseae",
    "LABEL_35": "Pneumonia",
    "LABEL_36": "Psoriasis",
    "LABEL_37": "Tuberculosis",
    "LABEL_38": "Typhoid",
    "LABEL_39": "Urinary tract infection",
    "LABEL_40": "Varicose veins"
}

def predict_disease(symptoms):
    # Get raw prediction (e.g., LABEL_23)
    results = pipe(symptoms)
    top_result = results[0]
    
    # Look up the human-readable name
    label_id = top_result['label'] # Gets "LABEL_23"
    disease_name = id2label.get(label_id, "Unknown Disease") # Converts to "Hepatitis E"
    
    return f"Predicted Disease: {disease_name}\nConfidence: {top_result['score']:.2f}"

# Create the interface
iface = gr.Interface(
    fn=predict_disease,
    inputs=gr.Textbox(lines=2, placeholder="I have skin rash and itching..."),
    outputs="text",
    title="AI Disease Symptom Checker",
    description="Enter your symptoms to get a prediction from the BERT model."
)

iface.launch()