import gradio as gr
import pandas as pd
import time
from collections import defaultdict

# Load the symptom-to-disease mapping
symptom_data = {
    "Shortness of breath": {
        "questions": [
            "Do you also have chest pain?",
            "Do you feel fatigued often?",
            "Have you noticed swelling in your legs?"
        ],
        "diseases": ["Atelectasis", "Emphysema", "Edema"],
        "weights_yes": [30, 30, 40],
        "weights_no": [10, 20, 30]
    },
    "Persistent cough": {
        "questions": [
            "Is your cough dry or with mucus?",
            "Do you experience fever?",
            "Do you have difficulty breathing?"
        ],
        "diseases": ["Pneumonia", "Fibrosis", "Infiltration"],
        "weights_yes": [35, 30, 35],
        "weights_no": [10, 15, 20]
    },
    "Sharp chest pain": {
        "questions": [
            "Does it worsen with deep breaths?",
            "Do you feel lightheaded?",
            "Have you had recent trauma or surgery?"
        ],
        "diseases": ["Pneumothorax", "Effusion", "Cardiomegaly"],
        "weights_yes": [40, 30, 30],
        "weights_no": [15, 20, 25]
    },
    "Fatigue & swelling": {
        "questions": [
            "Do you feel breathless when lying down?",
            "Have you gained weight suddenly?",
            "Do you experience irregular heartbeat?"
        ],
        "diseases": ["Edema", "Cardiomegaly"],
        "weights_yes": [50, 30, 20],
        "weights_no": [20, 15, 15]
    },
    "Chronic wheezing": {
        "questions": [
            "Do you have a history of smoking?",
            "Do you feel tightness in your chest?",
            "Do you have frequent lung infections?"
        ],
        "diseases": ["Emphysema", "Fibrosis"],
        "weights_yes": [40, 30, 30],
        "weights_no": [15, 25, 20]
    }
}

# Global variables to track user state
user_state = {}

def chatbot(user_input):
    if "state" not in user_state:
        user_state["state"] = "greet"
    
    if user_state["state"] == "greet":
        user_state["state"] = "ask_symptom"
        return "Hello! I'm a medical AI assistant. Please describe your primary symptom."
    
    elif user_state["state"] == "ask_symptom":
        if user_input not in symptom_data:
            return "I don't recognize that symptom. Please enter one of these: " + ", ".join(symptom_data.keys())
        user_state["symptom"] = user_input
        user_state["state"] = "ask_duration"
        return "How long have you been experiencing this symptom? (Less than a week / More than a week)"
    
    elif user_state["state"] == "ask_duration":
        if user_input.lower() == "less than a week":
            user_state.clear()
            return "It might be a temporary issue. Please monitor your symptoms and consult a doctor if they persist."
        elif user_input.lower() == "more than a week":
            user_state["state"] = "follow_up"
            user_state["current_question"] = 0
            user_state["disease_scores"] = defaultdict(int)
            return symptom_data[user_state["symptom"]]["questions"][0]
        else:
            return "Please respond with 'Less than a week' or 'More than a week'."
    
    elif user_state["state"] == "follow_up":
        symptom = user_state["symptom"]
        question_index = user_state["current_question"]
        
        # Update probabilities
        if user_input.lower() == "yes":
            for i, disease in enumerate(symptom_data[symptom]["diseases"]):
                user_state["disease_scores"][disease] += symptom_data[symptom]["weights_yes"][i]
        else:
            for i, disease in enumerate(symptom_data[symptom]["diseases"]):
                user_state["disease_scores"][disease] += symptom_data[symptom]["weights_no"][i]
        
        # Move to the next question or finish
        user_state["current_question"] += 1
        if user_state["current_question"] < len(symptom_data[symptom]["questions"]):
            return symptom_data[symptom]["questions"][user_state["current_question"]]
        
        # Final diagnosis
        probable_disease = max(user_state["disease_scores"], key=user_state["disease_scores"].get)
        user_state.clear()
        return f"Based on your symptoms, the most likely condition is: {probable_disease}. Please consult a doctor for confirmation."

# Gradio Chatbot UI with improved features
with gr.Blocks() as demo:
    gr.Markdown("# Conversational Image Recognition Assistant: AI-Powered X-ray Diagnosis for Healthcare")
    chatbot_ui = gr.Chatbot()
    user_input = gr.Textbox(placeholder="Enter your response...", label="Your Message")
    submit = gr.Button("Send")
    clear_chat = gr.Button("Clear Chat")
    
    def respond(user_message, history):
        history.append((user_message, "Thinking..."))  # Show thinking message
        yield history, ""  # Immediate update
        
        time.sleep(1.5)  # Simulate processing delay
        bot_response = chatbot(user_message)
        history[-1] = (user_message, bot_response)  # Update with real response
        yield history, ""
    
    submit.click(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
    user_input.submit(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
    clear_chat.click(lambda: ([], ""), outputs=[chatbot_ui, user_input])

demo.launch()