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 Key
groq_api_key = os.getenv("groq_api_key")
groq_client = Groq(api_key=groq_api_key)

# Load Zero-Shot 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")
}

# Translation 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)}"

# Function to generate expert 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. **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

# Function to create a medical report
def create_medical_report(symptoms):
    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_path = "medical_report.pdf"
    pdf.output(pdf_path)

    return full_report, pdf_path

# Gradio UI Setup
def show_section(selected):
    return (
        gr.update(visible=selected == "🏠 Home"),
        gr.update(visible=selected == "🔍 Diagnosis")
    )

with gr.Blocks() as interface:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 🏥 MedExpert")
            gr.Markdown("AI-powered medical diagnosis system.")
            nav_buttons = gr.Radio(["🏠 Home", "🔍 Diagnosis"], value="🏠 Home")

        with gr.Column(scale=4):
            with gr.Group(visible=True) as home_section:
                gr.Markdown("## 🌟 Welcome to MedExpert – Your AI-Powered Health Assistant 🏥🤖")
                gr.Markdown("""
    ### 🏆 Why Choose MedExpert?  
    - ✅ **Instant Symptom Analysis** 🤒  
    - ✅ **AI-Powered Medical Insights** 🧠  
    - ✅ **Expert Recommendations** 🩺  
    - ✅ **Generate and Download Medical Reports** 📄  
    """, elem_id="home-intro")
                gr.Markdown("---")
                gr.Markdown("""
    ### 💡 **How It Works?**
    1️⃣ **Describe your symptoms** ✍️  
    2️⃣ **AI analyzes & suggests a possible condition** 🧬  
    3️⃣ **Get expert analysis & a downloadable report** 📊  
    """)
                gr.Markdown("🔍 Click on **Diagnosis** above to start your health checkup!")

            with gr.Group(visible=False) as diagnostic_section:
                gr.Markdown("## 🩺 DiagnoGen AI")
                symptoms_input = gr.Textbox(label="Describe your symptoms", placeholder="e.g., headache, fever, nausea...")
                analyze_btn = gr.Button("Analyze Symptoms", variant="primary")
                report_output = gr.Textbox(label="Medical Report", interactive=False)
                pdf_output = gr.File(label="Download Full Report", file_count="single")

    # Actions
    analyze_btn.click(create_medical_report, inputs=symptoms_input, outputs=[report_output, pdf_output])
    nav_buttons.change(show_section, inputs=[nav_buttons], outputs=[home_section, diagnostic_section])

# Run Interface
interface.launch()