import gradio as gr
from transformers import pipeline

# Step 1: Load the AI Model
print("Loading the Health Analysis model...")
health_classifier = pipeline(
    "text-classification",
    model="shanover/disease_classifier_base"
)
print("Model loaded successfully!")

# Step 2: Create the Label-to-Name Dictionary
# This is the translator that maps the model's output to real disease names.
label_map = {
    'LABEL_0': 'Fungal infection', 'LABEL_1': 'Allergy', 'LABEL_2': 'GERD', 
    'LABEL_3': 'Chronic cholestasis', 'LABEL_4': 'Drug Reaction', 'LABEL_5': 'Peptic ulcer disease', 
    'LABEL_6': 'AIDS', 'LABEL_7': 'Diabetes', 'LABEL_8': 'Gastroenteritis', 
    'LABEL_9': 'Bronchial Asthma', 'LABEL_10': 'Hypertension', 'LABEL_11': 'Migraine', 
    'LABEL_12': 'Cervical spondylosis', 'LABEL_13': 'Paralysis (brain hemorrhage)', 
    'LABEL_14': 'Jaundice', 'LABEL_15': 'Malaria', 'LABEL_16': 'Chicken pox', 
    'LABEL_17': 'Dengue', 'LABEL_18': 'Typhoid', 'LABEL_19': 'Hepatitis A', 
    'LABEL_20': 'Hepatitis B', 'LABEL_21': 'Hepatitis C', 'LABEL_22': 'Hepatitis D', 
    'LABEL_23': 'Hepatitis E', 'LABEL_24': 'Alcoholic hepatitis', 'LABEL_25': 'Tuberculosis', 
    'LABEL_26': 'Common Cold', 'LABEL_27': 'Pneumonia', 'LABEL_28': 'Dimorphic hemmorhoids(piles)', 
    'LABEL_29': 'Heart attack', 'LABEL_30': 'Varicose veins', 'LABEL_31': 'Hypothyroidism', 
    'LABEL_32': 'Hyperthyroidism', 'LABEL_33': 'Hypoglycemia', 'LABEL_34': 'Osteoarthristis', 
    'LABEL_35': 'Arthritis', 'LABEL_36': '(vertigo) Paroymsal Positional Vertigo', 
    'LABEL_37': 'Acne', 'LABEL_38': 'Urinary tract infection', 'LABEL_39': 'Psoriasis', 
    'LABEL_40': 'Impetigo'
}

# Step 3: Define the main analysis function
def analyze_symptoms(symptoms):
    # Get the prediction dictionary from the model, e.g., {'label': 'LABEL_2', 'score': 0.96}
    prediction_dict = health_classifier(symptoms)[0]
    
    # Get the generic label (e.g., 'LABEL_2')
    generic_label = prediction_dict['label']
    
    # Look up the real disease name in our dictionary
    disease = label_map.get(generic_label, "Unknown Condition")
    
    # Get the confidence score
    confidence_score = prediction_dict['score']
    
    # Format the final, clean string for the user
    return f"Predicted Condition: {disease}\nConfidence Score: {confidence_score:.2f}"

# Step 4: Define the app's text content
app_title = "Symptom to Disease Classifier"
app_description = "Enter symptoms to predict a possible related medical condition. For educational use only."
disclaimer = "⚠️ **IMPORTANT DISCLAIMER:** This is NOT a medical diagnostic tool. Always consult a qualified healthcare professional for medical advice."

# Step 5: Create and Launch the Gradio Interface
app = gr.Interface(
    fn=analyze_symptoms,
    inputs=gr.Textbox(
        lines=4,
        placeholder="e.g., 'I am feeling vomiting, breathlessness, and sweating...'"
    ),
    outputs="text",
    title=app_title,
    description=app_description,
    article=disclaimer,
    examples=[
        ["I have a persistent cough, high fever, and body aches."],
        ["Feeling very thirsty and needing to urinate frequently."],
        ["I have itchy skin with red patches and silver scales."]
    ]
)

app.launch()