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

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

# ✅ Use Google Translate for auto-detection & translation
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)}"

# ✅ Load Models
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")
}

def generate_expert_analysis(condition, symptoms):
    """Generates expert medical analysis using Groq API"""
    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. **No AI disclaimers or generic advice**.
    """

    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

def create_medical_report(symptoms):
    """Generates a complete medical report"""
    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\n"
            "**Key Questions for Your Doctor:**\n"
            "1. Is this condition acute or chronic?\n"
            "2. What medication options are suitable?\n"
            "3. What lifestyle changes help manage this condition?\n"
            "4. What warning signs require immediate attention?\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

# ✅ Create Sidebar with Navigation
with gr.Blocks(css="body { background-color: #f5f7fa; }") as app:
    with gr.Row():
        gr.Markdown("<h1 style='text-align: center; color: #4A90E2;'>MedExpert AI</h1>")
    
    with gr.Tabs():
        with gr.TabItem("🏠 Home"):
            gr.Markdown("""
                <h2 style="color: #4A90E2;">Welcome to MedExpert</h2>
                <p style="font-size:18px;">An AI-powered medical assistant that provides expert analysis based on your symptoms.</p>
                <p style="font-size:18px;">Simply describe your symptoms, and let AI generate a detailed report!</p>
            """)

        with gr.TabItem("📄 Medical Diagnosis"):
            with gr.Row():
                symptoms_input = gr.Textbox(label="Describe Your Symptoms", placeholder="e.g., Mujhay sar main dard hai", interactive=True)
                report_btn = gr.Button("🔍 Generate Report", elem_id="generate-btn")
            
            with gr.Row():
                report_output = gr.Textbox(label="Complete Medical Report", interactive=False)
                pdf_output = gr.File(label="Download PDF Report")

            report_btn.click(create_medical_report, inputs=[symptoms_input], outputs=[report_output, pdf_output])

        with gr.TabItem("ℹ️ About"):
            gr.Markdown("""
                <h2 style="color: #4A90E2;">About This App</h2>
                <p style="font-size:18px;">MedExpert AI is an intelligent health assistant designed to provide expert insights into medical conditions.</p>
            """)

# ✅ Launch App
app.launch()