import gradio as gr
from transformers import pipeline

# Use a zero-shot classification model
model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Candidate labels for medical diagnosis
candidate_labels = [
    "respiratory infection", "viral infection", "bacterial infection", "autoimmune disease", 
    "cardiovascular disease", "endocrine disorders", "gastrointestinal disorders", 
    "neurological disorders", "skin disorders", "genetic disorders", "cancer", "kidney disease"
]

# Function to diagnose based on symptoms
def diagnose(symptoms):
    result = model(symptoms, candidate_labels=candidate_labels)
    diagnosis_category = result['labels'][0]  # Top predicted category
    confidence = result['scores'][0]  # Confidence score

    # Only accept predictions with confidence >= 0.60
    if confidence < 0.60:
        return "Unable to provide a confident diagnosis. Please consult a healthcare professional."

    return f"Predicted diagnosis category: {diagnosis_category} with confidence: {confidence:.2f}"

# Triage function to assess symptom severity
def triage(symptoms):
    # Check for critical symptoms that need urgent attention
    if "shortness of breath" in symptoms or "chest pain" in symptoms:
        return "Urgent: Seek immediate medical attention."

    # Check for milder symptoms
    elif "fever" in symptoms and "cough" in symptoms:
        return "Mild: Likely a viral infection, monitor symptoms."

    # New logic for symptoms like "red nose and cheeks"
    elif "red nose" in symptoms or "red cheeks" in symptoms:
        return "Non-urgent: Likely a mild skin reaction. Monitor for any additional symptoms."

    # Default case if no specific symptoms are detected
    else:
        return "Non-urgent: No immediate concern."

# Combine diagnosis and triage into one function
def full_check(symptoms):
    diagnosis = diagnose(symptoms)  # Get diagnosis category
    severity = triage(symptoms)     # Get severity level
    return diagnosis, severity

# Create the Gradio interface
iface = gr.Interface(
    fn=full_check, 
    inputs="text", 
    outputs=["text", "text"], 
    title="Sehat Guard - Symptom Checker",
    description="Enter your symptoms to get a possible diagnosis and severity of the condition."
)

# Launch the app
iface.launch()