app.py

import gradio as gr
from transformers import pipeline
from fpdf import FPDF
import os
from groq import Groq
from deep_translator import GoogleTranslator

# ✅ Load API Keys
groq_api_key = os.getenv("groq_api_key")  
groq_client = Groq(api_key=groq_api_key)

# ✅ Load Zero-Shot Disease Classification Model
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# ✅ Medical Conditions List
conditions = [
    "Asthma", "COPD", "Pneumonia", "Tuberculosis", "COVID-19", "Bronchitis",
    "Heart Failure", "Hypertension", "Diabetes Type 1", "Diabetes Type 2",
    "Migraine", "Gastroenteritis", "Anemia", "Depression", "Anxiety Disorder",
    "Chronic Kidney Disease", "UTI", "Osteoporosis", "Psoriasis", "Epilepsy"
]

# ✅ Specialist Mapping
specialist_mapping = {
    "Asthma": ("Pulmonologist", "Respiratory System"),
    "COPD": ("Pulmonologist", "Respiratory System"),
    "Pneumonia": ("Pulmonologist", "Respiratory System"),
    "Tuberculosis": ("Infectious Disease Specialist", "Respiratory System"),
    "COVID-19": ("Infectious Disease Specialist", "Immune System"),
    "Heart Failure": ("Cardiologist", "Cardiovascular System"),
    "Hypertension": ("Cardiologist", "Cardiovascular System"),
    "Diabetes Type 1": ("Endocrinologist", "Endocrine System"),
    "Diabetes Type 2": ("Endocrinologist", "Endocrine System"),
    "Migraine": ("Neurologist", "Nervous System"),
    "Gastroenteritis": ("Gastroenterologist", "Digestive System"),
    "Anemia": ("Hematologist", "Blood Disorders"),
    "Depression": ("Psychiatrist", "Mental Health"),
    "Anxiety Disorder": ("Psychiatrist", "Mental Health"),
    "Chronic Kidney Disease": ("Nephrologist", "Urinary System"),
    "UTI": ("Urologist", "Urinary System"),
    "Osteoporosis": ("Orthopedic Specialist", "Musculoskeletal System"),
    "Psoriasis": ("Dermatologist", "Skin Disorders"),
    "Epilepsy": ("Neurologist", "Nervous System")
}

# ✅ Translate Function
def translate_text(text, target_lang="en"):
    try:
        return GoogleTranslator(source="auto", target=target_lang).translate(text)
    except Exception as e:
        return f"Translation Error: {str(e)}"

# ✅ Expert AI Analysis
def generate_expert_analysis(condition, symptoms):
    specialist_title = specialist_mapping.get(condition, ("General Physician", "General Medicine"))[0]
    prompt = f"""As a {specialist_title.lower()}, explain {condition} to a patient experiencing these symptoms: "{symptoms}".
    Structure the response into:
    1. **Biological Process**  
    2. **Immediate Treatment**  
    3. **Long-term Care**  
    4. **Emergency Signs**  
    5. **Diet Plan**  
    Use professional yet simple language.
    """
    response = groq_client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="mixtral-8x7b-32768",
        temperature=0.5,
        max_tokens=1024
    )
    return response.choices[0].message.content

# ✅ Medical Report Generator
def create_medical_report(symptoms):
    try:
        translated_symptoms = translate_text(symptoms, "en")
        result = classifier(translated_symptoms, conditions, multi_label=False)
        diagnosis = result['labels'][0]  
        specialist, system = specialist_mapping.get(diagnosis, ("General Physician", "General Medicine"))
        expert_analysis = generate_expert_analysis(diagnosis, translated_symptoms)

        full_report = (
            f"**Medical Report**\n\n"
            f"**Patient Symptoms:** {translated_symptoms}\n"
            f"**Primary Diagnosis:** {diagnosis}\n"
            f"**Affected System:** {system}\n"
            f"**Consult:** {specialist}\n\n"
            f"**Expert Analysis:**\n{expert_analysis}\n"
        )

        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", size=12)
        pdf.multi_cell(0, 10, full_report)
        pdf.output("report.pdf")

        return full_report, "report.pdf"

    except Exception as e:
        return f"Error generating report: {str(e)}", None

# ✅ Medical Chatbot Function
def medical_chatbot(user_input, chat_history):
    system_prompt = """You are a compassionate psychiatric assistant named MedMind. Your role is to:
    - Provide emotional support and stress relief techniques
    - Offer evidence-based mental health advice
    - Help users understand medical terminology
    - Never diagnose but suggest professional help when needed
    - Maintain therapeutic conversation flow
    - Remember previous interactions for context"""
    
    messages = [
        {"role": "system", "content": system_prompt},
        *[{"role": "user" if i%2==0 else "assistant", "content": msg} 
          for i, msg in enumerate(chat_history)]
    ]
    
    if user_input.strip():
        messages.append({"role": "user", "content": user_input})
        
        response = groq_client.chat.completions.create(
            messages=messages,
            model="mixtral-8x7b-32768",
            temperature=0.7,
            max_tokens=500
        )
        
        bot_response = response.choices[0].message.content
        chat_history.append((user_input, bot_response))
        
    return "", chat_history

# ✅ Gradio UI with Sidebar
css = """
body { background-color: #f0f4f8; }
.sidebar { background-color: #2c3e50; padding: 20px; height: 100vh; color: white; }
.chat-container { height: 70vh; overflow-y: auto; }
.message { padding: 10px; margin: 5px; border-radius: 5px; }
.user-message { background-color: #e3f2fd; margin-left: 20%; }
.bot-message { background-color: #ffffff; margin-right: 20%; }
"""

with gr.Blocks(css=css) as interface:
    chat_history = gr.State([])
    
    with gr.Row():
        with gr.Column(scale=1, elem_classes="sidebar"):
            gr.Markdown("## MedExpert AI")
            with gr.Group():
                gr.Button("🏠 Home", variant="secondary")
                gr.Button("🩺 Diagnostic", variant="secondary")
                gr.Button("💬 Mental Health Chat", variant="secondary")
        
        with gr.Column(scale=4):
            # Home Section
            with gr.Group(visible=True) as home_section:
                gr.Markdown("# Welcome to MedExpert AI")
                gr.Markdown("Your intelligent medical assistant providing:")
                gr.Markdown("- Symptom Analysis\n- Medical Reports\n- Mental Health Support")
            
            # Diagnostic Section
            with gr.Group(visible=False) as diagnostic_section:
                symptoms_input = gr.Textbox(label="Describe Your Symptoms")
                analyze_btn = gr.Button("Analyze Symptoms 🏥", variant="primary")
                report_output = gr.Textbox(label="Medical Report", interactive=False)
                pdf_output = gr.File(label="Download PDF Report")
            
            # Chatbot Section
            with gr.Group(visible=False) as chatbot_section:
                chatbot = gr.Chatbot(elem_classes="chat-container")
                user_input = gr.Textbox(placeholder="Type your message here...", show_label=False)
                send_btn = gr.Button("Send", variant="primary")
    
    # Diagnostic Logic
    analyze_btn.click(
        create_medical_report,
        inputs=[symptoms_input],
        outputs=[report_output, pdf_output]
    )
    
    # Chatbot Logic
    send_btn.click(
        medical_chatbot,
        inputs=[user_input, chat_history],
        outputs=[user_input, chatbot]
    )
    user_input.submit(
        medical_chatbot,
        inputs=[user_input, chat_history],
        outputs=[user_input, chatbot]
    )

interface.launch()