import gradio as gr
import time
from threading import Thread

# ===========================
# Malaria Symptom Rules
# ===========================
malaria_rules = {
    "fever": {"diagnosis":["Patient may have malaria."],"treatment":["Recommend antimalarial therapy."],"prognosis":["Mild if treated early."]},
    "chills": {"diagnosis":["Chills often accompany malaria fever."],"treatment":["Supportive care, hydration, antipyretics."],"prognosis":["Typically resolves with treatment."]},
    "headache": {"diagnosis":["Common in malaria patients."],"treatment":["Analgesics as needed."],"prognosis":["Improves with treatment."]},
    "nausea": {"diagnosis":["May indicate malaria or gastrointestinal infection."],"treatment":["Anti-nausea medications if severe."],"prognosis":["Resolves with malaria treatment."]},
    "vomiting": {"diagnosis":["May indicate severe malaria if combined with fever."],"treatment":["Consider oral or IV hydration."],"prognosis":["Monitor for dehydration."]},
    "muscle_aches": {"diagnosis":["Common symptom of malaria."],"treatment":["Analgesics if necessary."],"prognosis":["Improves with antimalarial therapy."]},
    "fatigue": {"diagnosis":["Common symptom, may persist during recovery."],"treatment":["Rest and supportive care."],"prognosis":["Gradual improvement expected."]},
    "sweating": {"diagnosis":["Sweating often follows fever spikes."],"treatment":["Supportive care."],"prognosis":["Normalizes with treatment."]},
    "anemia": {"diagnosis":["Severe malaria can cause hemolytic anemia."],"treatment":["Monitor hemoglobin; blood transfusion if critical."],"prognosis":["Serious if untreated."]},
    "jaundice": {"diagnosis":["May indicate liver involvement; risk for severe malaria."],"treatment":["Monitor liver function, supportive care."],"prognosis":["Can be severe if untreated."]},
    "seizures": {"diagnosis":["Cerebral malaria, medical emergency."],"treatment":["Immediate hospitalization, IV antimalarials."],"prognosis":["Life-threatening without prompt care."]},
    # Combined symptom rules
    ("fever","chills"): {"diagnosis":["Classic malaria presentation."],"treatment":["Prompt antimalarial therapy."],"prognosis":["Good if treated early."]},
    ("fever","vomiting"): {"diagnosis":["Possible severe malaria."],"treatment":["Monitor hydration, antimalarials."],"prognosis":["Serious if untreated."]},
    ("jaundice","anemia"): {"diagnosis":["Severe malaria with liver involvement."],"treatment":["Hospitalization recommended."],"prognosis":["Severe, close monitoring required."]},
    ("seizures","fever"): {"diagnosis":["Cerebral malaria; emergency."],"treatment":["Immediate IV therapy and ICU support."],"prognosis":["Life-threatening."]}
}

# ===========================
# Analysis Functions
# ===========================
def analyze_malaria(symptoms):
    diagnosis = set()
    treatment = set()
    prognosis = set()
    # Combined symptoms first
    for combo in [k for k in malaria_rules if isinstance(k, tuple)]:
        if all(sym in symptoms for sym in combo):
            diagnosis.update(malaria_rules[combo]["diagnosis"])
            treatment.update(malaria_rules[combo]["treatment"])
            prognosis.update(malaria_rules[combo]["prognosis"])
    # Single symptoms
    for sym in symptoms:
        if sym in malaria_rules:
            diagnosis.update(malaria_rules[sym]["diagnosis"])
            treatment.update(malaria_rules[sym]["treatment"])
            prognosis.update(malaria_rules[sym]["prognosis"])
    return {
        "Diagnosis": "\n".join(diagnosis) if diagnosis else "No malaria-related diagnosis detected.",
        "Treatment": "\n".join(treatment) if treatment else "No treatment recommended.",
        "Prognosis": "\n".join(prognosis) if prognosis else "No prognosis available."
    }

def run_analysis_with_loading(age, gender, location, travel, travel_details, symptoms_list,
                              temperature, blood_pressure, heart_rate, previous_malaria, medications, additional_notes):
    # Show loading dots first
    loading_html = "<div style='text-align:center; font-size:1.5rem;'><div>🤖 AI is analyzing patient data<span id='dots'>...</span></div></div>"
    diag_output = treat_output = prog_output = loading_html
    
    # Run the actual analysis in a thread to simulate loading
    def compute():
        nonlocal diag_output, treat_output, prog_output
        time.sleep(2)  # simulate AI processing time
        symptoms = symptoms_list
        result = analyze_malaria(symptoms)
        severity = "Mild"
        color = "#10B981"
        if all(sym in symptoms for sym in ["seizures","fever"]):
            severity = "Life-threatening"
            color = "#EF4444"
        elif all(sym in symptoms for sym in ["jaundice","anemia"]):
            severity = "Severe"
            color = "#F97316"
        
        alert_banner = f"<div style='border-radius:0.5rem; padding:1rem; background-color:{color}; color:white; font-weight:bold; margin-bottom:1rem;'>⚠️ Severity Alert: {severity}</div>"
        diag_output = f"{alert_banner}<b>Diagnosis:</b><br>{result['Diagnosis'].replace(chr(10),'<br>')}"
        treat_output = f"{alert_banner}<b>Treatment:</b><br>{result['Treatment'].replace(chr(10),'<br>')}"
        prog_output = f"{alert_banner}<b>Prognosis:</b><br>{result['Prognosis'].replace(chr(10),'<br>')}"
    
    thread = Thread(target=compute)
    thread.start()
    thread.join()
    
    return diag_output, treat_output, prog_output

# ===========================
# Gradio UI
# ===========================
symptoms_options = ["fever","chills","headache","nausea","vomiting","muscle_aches","fatigue","sweating","anemia","jaundice","seizures"]

with gr.Blocks(css="""
body {background-color:#F9FAFB; font-family:sans-serif;}
.header {background:linear-gradient(135deg,#667eea 0%,#764ba2 100%); color:white; padding:2rem; text-align:center; border-radius:1rem; margin-bottom:1rem;}
.card {border-radius:1rem; box-shadow:0 10px 25px rgba(0,0,0,0.1); background:white; padding:1.5rem; margin-bottom:1rem;}
.btn {background-color:#5D5CDE; color:white; padding:0.75rem 1.5rem; border-radius:0.75rem; font-weight:bold; cursor:pointer;}
.btn:hover {background-color:#4B4AC8;}
.checkbox-group label {display:block; padding:0.5rem; border-radius:0.5rem; cursor:pointer; transition:0.2s;}
.checkbox-group label:hover {background-color:#E0E7FF;}
""") as demo:
    
    gr.HTML("<div class='header'><h1 class='text-4xl font-bold'>🦟 Malaria AI Assistant</h1><p>Advanced diagnostic, treatment, and prognostic analysis</p></div>")
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(elem_classes="card"):
                age = gr.Number(label="Age")
                gender = gr.Dropdown(["male","female","other"], label="Gender")
                location = gr.Textbox(label="Location")
                travel = gr.Checkbox(label="Recent Travel to Endemic Areas")
                travel_details = gr.Textbox(label="Travel Details")
                symptoms_list = gr.CheckboxGroup(symptoms_options, label="Symptoms", elem_classes="checkbox-group")
                temperature = gr.Number(label="Temperature (°C)")
                blood_pressure = gr.Textbox(label="Blood Pressure")
                heart_rate = gr.Number(label="Heart Rate")
                previous_malaria = gr.Checkbox(label="Previous Malaria Episodes")
                medications = gr.Textbox(label="Current Medications / Allergies")
                additional_notes = gr.Textbox(label="Additional Notes")
        
        with gr.Column(scale=1):
            with gr.Group(elem_classes="card"):
                diag_output = gr.HTML(label="Diagnosis", value="<div style='text-align:center;'>🤖 Waiting for input...</div>")
                treat_output = gr.HTML(label="Treatment", value="<div style='text-align:center;'>🤖 Waiting for input...</div>")
                prog_output = gr.HTML(label="Prognosis", value="<div style='text-align:center;'>🤖 Waiting for input...</div>")
    
    run_btn = gr.Button("Run Analysis", elem_classes="btn")
    run_btn.click(
        run_analysis_with_loading,
        inputs=[age, gender, location, travel, travel_details, symptoms_list,
                temperature, blood_pressure, heart_rate, previous_malaria,
                medications, additional_notes],
        outputs=[diag_output, treat_output, prog_output]
    )

if __name__ == "__main__":
    demo.launch()