import gradio as gr
from transformers import pipeline
import requests

class AIHealthAssistant:
    def __init__(self):
        # Initialize symptom checker model
        self.symptom_checker = pipeline(
            "text-classification",
            model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
        )
        
        # Initialize medical QA model
        self.medical_qa = pipeline(
            "question-answering",
            model="deepset/roberta-base-squad2"
        )
        
        # Knowledge base (can be replaced with API calls)
        self.disease_db = {
            "influenza": {
                "symptoms": ["fever", "cough", "sore throat", "runny nose", "body aches"],
                "advice": "Rest, stay hydrated, take fever reducers like acetaminophen",
                "precautions": ["Annual flu vaccine", "Frequent hand washing", "Avoid close contact"]
            },
            "migraine": {
                "symptoms": ["severe headache", "nausea", "sensitivity to light", "aura"],
                "advice": "Rest in dark room, take prescribed medication, apply cold compress",
                "precautions": ["Identify triggers", "Maintain sleep schedule", "Stay hydrated"]
            }
        }
    
    def get_disease_from_symptoms(self, symptoms):
        """Predict disease from symptoms using Hugging Face model"""
        try:
            # For production, replace with a proper medical model
            result = self.symptom_checker(symptoms)
            
            # Map to diseases in our database
            for disease in self.disease_db:
                if disease in symptoms.lower():
                    return disease
            
            # Fallback to first disease (in real app, use proper mapping)
            return list(self.disease_db.keys())[0]
        
        except Exception as e:
            print(f"Model error: {e}")
            return "unknown"
    
    def get_medical_info(self, disease, info_type):
        """Get medical information from database or API"""
        disease = disease.lower()
        
        # Check local database first
        if disease in self.disease_db:
            return self.disease_db[disease].get(info_type, "Information not available")
        
        # Fallback to API (example using hypothetical medical API)
        try:
            if info_type == "symptoms":
                prompt = f"What are the symptoms of {disease}?"
            elif info_type == "advice":
                prompt = f"What is the treatment for {disease}?"
            else:
                prompt = f"What precautions should be taken for {disease}?"
            
            # In a real app, replace with actual API call:
            # response = requests.get(f"https://medical-api.example.com/{disease}")
            # return response.json().get(info_type)
            
            # For demo, using the QA model
            context = f"{disease} is a medical condition. {self.get_medical_advice_from_api(disease)}"
            result = self.medical_qa(question=prompt, context=context)
            return result['answer']
        
        except Exception as e:
            print(f"API error: {e}")
            return "Information not available"
    
    def get_medical_advice_from_api(self, disease):
        """Simulate API response for demo purposes"""
        api_responses = {
            "diabetes": "Diabetes requires blood sugar monitoring, insulin therapy, and dietary changes.",
            "hypertension": "Hypertension management includes medication, low-salt diet, and regular exercise."
        }
        return api_responses.get(disease.lower(), "Consult a healthcare professional for proper diagnosis and treatment.")

def create_demo():
    assistant = AIHealthAssistant()
    
    def process_input(user_input, mode):
        if mode == "Symptoms to Disease":
            disease = assistant.get_disease_from_symptoms(user_input)
            symptoms = assistant.get_medical_info(disease, "symptoms")
            advice = assistant.get_medical_info(disease, "advice")
            precautions = assistant.get_medical_info(disease, "precautions")
            
            output = (
                f"🔍 Possible Condition: {disease.capitalize()}\n\n"
                f"📋 Symptoms:\n- " + "\n- ".join(symptoms) + "\n\n"
                f"💊 Recommended Actions:\n{advice}\n\n"
                f"🛡️ Precautions:\n- " + "\n- ".join(precautions)
            )
        
        elif mode == "Disease to Symptoms":
            symptoms = assistant.get_medical_info(user_input, "symptoms")
            advice = assistant.get_medical_info(user_input, "advice")
            
            output = (
                f"📋 Symptoms of {user_input.capitalize()}:\n- " + "\n- ".join(symptoms) + "\n\n"
                f"💊 Recommended Actions:\n{advice}"
            )
        
        return output
    
    with gr.Blocks(title="AI Health Assistant") as demo:
        gr.Markdown("# 🏥 AI Health Assistant")
        gr.Markdown("Enter symptoms or a disease name to get medical information")
        
        with gr.Row():
            input_mode = gr.Radio(
                choices=["Symptoms to Disease", "Disease to Symptoms"],
                label="Input Mode"
            )
            user_input = gr.Textbox(
                label="Input",
                placeholder="Enter symptoms or disease name..."
            )
        
        submit_btn = gr.Button("Get Medical Information")
        output = gr.Textbox(label="Result", interactive=False, lines=10)
        
        # Example inputs
        gr.Examples(
            examples=[
                ["headache, nausea, sensitivity to light", "Symptoms to Disease"],
                ["influenza", "Disease to Symptoms"],
                ["fever, cough, sore throat", "Symptoms to Disease"]
            ],
            inputs=[user_input, input_mode],
            outputs=output,
            fn=process_input,
            cache_examples=True
        )
        
        submit_btn.click(
            fn=process_input,
            inputs=[user_input, input_mode],
            outputs=output
        )
        
        gr.Markdown("""
        ## ⚠️ Important Disclaimer
        This AI assistant provides general health information only and is not a substitute 
        for professional medical advice, diagnosis, or treatment.
        Always seek the advice of your physician or other qualified health provider 
        with any questions you may have regarding a medical condition.
        """)
    
    return demo

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