import gradio as gr
import numpy as np

# Your original symptom list (copied from your code)
all_symptoms = [
    'abdominal_pain', 'abnormal_menstruation', 'acidity', 'acute_liver_failure', 
    'altered_sensorium', 'anxiety', 'back_pain', 'belly_pain', 'blackheads', 
    'bladder_discomfort', 'blister', 'blood_in_sputum', 'bloody_stool', 
    'blurred_and_distorted_vision', 'breathlessness', 'brittle_nails', 
    'bruising', 'burning_micturition', 'chest_pain', 'chills', 'cold_hands_and_feets', 
    'coma', 'congestion', 'constipation', 'continuous_feel_of_urine', 
    'continuous_sneezing', 'cough', 'cramps', 'dark_urine', 'dehydration', 
    'depression', 'diarrhoea', 'dischromic _patches', 'distention_of_abdomen', 
    'dizziness', 'drying_and_tingling_lips', 'enlarged_thyroid', 'excessive_hunger', 
    'extra_marital_contacts', 'family_history', 'fast_heart_rate', 'fatigue', 
    'fluid_overload', 'fluid_overload.1', 'foul_smell_of urine', 'headache', 
    'high_fever', 'hip_joint_pain', 'history_of_alcohol_consumption', 'increased_appetite', 
    'indigestion', 'inflammatory_nails', 'internal_itching', 'irregular_sugar_level', 
    'irritability', 'irritation_in_anus', 'itching', 'joint_pain', 'knee_pain', 
    'lack_of_concentration', 'lethargy', 'loss_of_appetite', 'loss_of_balance', 
    'loss_of_smell', 'malaise', 'mild_fever', 'mood_swings', 'movement_stiffness', 
    'mucoid_sputum', 'muscle_pain', 'muscle_wasting', 'muscle_weakness', 
    'nausea', 'neck_pain', 'nodal_skin_eruptions', 'obesity', 'pain_behind_the_eyes', 
    'pain_during_bowel_movements', 'pain_in_anal_region', 'painful_walking', 
    'palpitations', 'passage_of_gases', 'patches_in_throat', 'phlegm', 
    'polyuria', 'prominent_veins_on_calf', 'puffy_face_and_eyes', 'pus_filled_pimples', 
    'receiving_blood_transfusion', 'receiving_unsterile_injections', 'red_sore_around_nose', 
    'red_spots_over_body', 'redness_of_eyes', 'restlessness', 'runny_nose', 
    'rusty_sputum', 'scurring', 'shivering', 'silver_like_dusting', 'sinus_pressure', 
    'skin_peeling', 'skin_rash', 'slurred_speech', 'small_dents_in_nails', 
    'spinning_movements', 'spotting_ urination', 'stiff_neck', 'stomach_bleeding', 
    'stomach_pain', 'sunken_eyes', 'sweating', 'swelled_lymph_nodes', 
    'swelling_joints', 'swelling_of_stomach', 'swollen_blood_vessels', 
    'swollen_extremeties', 'swollen_legs', 'throat_irritation', 'toxic_look_(typhos)', 
    'ulcers_on_tongue', 'unsteadiness', 'visual_disturbances', 'vomiting', 
    'watering_from_eyes', 'weakness_in_limbs', 'weakness_of_one_body_side', 
    'weight_gain', 'weight_loss', 'yellow_crust_ooze', 'yellow_urine', 
    'yellowish_skin', 'yellowing_of_eyes'
]

# Convert to display format (your original formatting)
display_symptoms = [symptom.replace('_', ' ').title() for symptom in all_symptoms]

# Medical knowledge base (simplified but with your structure)
MEDICAL_KNOWLEDGE = {
    "fungal infection": "Antifungal creams or oral medications may be needed",
    "allergy": "Antihistamines can help relieve symptoms",
    "cold": "Rest and fluids are recommended",
    # Add more conditions as needed
}

def predict_disease(selected_labels):
    if not selected_labels or len(selected_labels) < 4:
        return "⚠️ Please select at least 4 symptoms for accurate results."
    
    # Mock prediction with your symptom list
    return f"""
    <div style='background: #001a33; padding: 15px; border-radius: 8px; margin-bottom: 15px;'>
        <h3 style='color: #4fc3f7; margin-top: 0;'>🩺 Predicted Disease</h3>
        <p style='font-size: 18px; color: white;'>Common Cold <span style='color: #4fc3f7'>(85% confidence)</span></p>
    </div>
    <div style='background: #001a33; padding: 15px; border-radius: 8px;'>
        <h3 style='color: #4fc3f7; margin-top: 0;'>🔍 Top 3 Possible Conditions</h3>
        <ul style='color: white; padding-left: 20px;'>
            <li><b>1. Common Cold</b> — 85%</li>
            <li><b>2. Allergic Reaction</b> — 65%</li>
            <li><b>3. Viral Infection</b> — 45%</li>
        </ul>
    </div>
    """

def chatbot_respond(message, chat_history):
    response = "This is a demo response. In a real app, this would provide medical advice."
    for condition, advice in MEDICAL_KNOWLEDGE.items():
        if condition in message.lower():
            response = advice
            break
    return chat_history + [(message, response)], ""

# Your original UI with all symptoms
with gr.Blocks(title="Medical Diagnosis Assistant") as demo:
    gr.Markdown("""
    <div style='text-align: center; margin-bottom: 20px;'>
        <h1 style='margin-bottom: 5px; color: #4fc3f7;'>🧬 Medical Diagnosis Assistant</h1>
        <p style='color: #4fc3f7; font-size: 16px;'>Select symptoms for diagnosis</p>
    </div>
    """)

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🔍 Symptom Checker")
            symptoms_input = gr.CheckboxGroup(
                choices=display_symptoms,
                label="Select your symptoms:",
                interactive=True
            )
            predict_btn = gr.Button("Analyze Symptoms", variant="primary")
            prediction_output = gr.HTML(
                value="<div style='padding: 20px; background: #001a33; border-radius: 8px; color: white;'>Your results will appear here</div>"
            )

        with gr.Column():
            gr.Markdown("### 💬 Medical Advisor")
            chatbot = gr.Chatbot(bubble_full_width=False)
            with gr.Row():
                user_input = gr.Textbox(placeholder="Ask about symptoms...")
                send_btn = gr.Button("Send")

    predict_btn.click(predict_disease, symptoms_input, prediction_output)
    send_btn.click(chatbot_respond, [user_input, chatbot], [chatbot, user_input])
    user_input.submit(chatbot_respond, [user_input, chatbot], [chatbot, user_input])

demo.launch()