app.py

import gradio as gr
from transformers import pipeline
from fpdf import FPDF
import re
from groq import Groq  # Groq API
import os

groq_api_key = os.environ["groq_api_key"] 

# Initialize Groq client
groq_client = Groq(api_key=groq_api_key)  # Replace with your Groq API key

# Load models
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-hi-en")
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Expanded Medical Conditions
conditions = [
    # Respiratory Diseases
    "Asthma", "COPD", "Pneumonia", "Tuberculosis", "COVID-19", "Bronchitis", "Pulmonary Fibrosis", "Cystic Fibrosis", "Emphysema", "Sleep Apnea",
    # Cardiovascular Diseases
    "Heart Failure", "Hypertension", "Hypotension", "Coronary Artery Disease", "Arrhythmia", "Atrial Fibrillation", "Myocardial Infarction", "Peripheral Artery Disease", "Stroke", "Aneurysm",
    # Endocrine & Metabolic Diseases
    "Diabetes Type 1", "Diabetes Type 2", "Hyperthyroidism", "Hypothyroidism", "Addison’s Disease", "Cushing’s Syndrome", "Polycystic Ovary Syndrome", "Metabolic Syndrome", "Hypoglycemia", "Hyperglycemia",
    # Neurological Disorders
    "Migraine", "Epilepsy", "Parkinson’s Disease", "Multiple Sclerosis", "Alzheimer’s Disease", "Huntington’s Disease", "ALS", "Meningitis", "Encephalitis", "Bell’s Palsy",
    # Digestive System Diseases
    "Gastroenteritis", "Acid Reflux", "Peptic Ulcer", "Crohn’s Disease", "Ulcerative Colitis", "Irritable Bowel Syndrome", "Hepatitis A", "Hepatitis B", "Hepatitis C", "Cirrhosis", "Gallstones", "Pancreatitis",
    # Blood & Immune System Disorders
    "Anemia", "Leukemia", "Lymphoma", "Sickle Cell Disease", "Hemophilia", "Thrombocytopenia", "Lupus", "Rheumatoid Arthritis", "HIV/AIDS", "Sepsis", "Aplastic Anemia",
    # Musculoskeletal Disorders
    "Arthritis", "Osteoporosis", "Gout", "Fibromyalgia", "Scoliosis", "Ankylosing Spondylitis", "Carpal Tunnel Syndrome", "Muscular Dystrophy", "Tendinitis", "Bone Fractures",
    # Skin Conditions
    "Eczema", "Psoriasis", "Acne", "Rosacea", "Fungal Infections", "Vitiligo", "Hives", "Melanoma", "Alopecia", "Sunburn",
    # Infectious Diseases
    "Influenza", "Malaria", "Dengue Fever", "Chickenpox", "Measles", "Mumps", "Typhoid Fever", "Tetanus", "Rabies", "Lyme Disease",
    # Kidney & Urinary Tract Diseases
    "Chronic Kidney Disease", "Kidney Stones", "UTI", "Bladder Cancer", "Polycystic Kidney Disease", "Nephritis", "Hydronephrosis", "Glomerulonephritis", "Interstitial Cystitis", "Prostatitis",
    # Reproductive System Disorders
    "Endometriosis", "Ovarian Cysts", "Erectile Dysfunction", "Prostate Cancer", "Uterine Fibroids", "Pelvic Inflammatory Disease", "Infertility", "Testicular Cancer", "Menstrual Disorders", "Premature Ovarian Failure",
    # Mental Health Disorders
    "Depression", "Anxiety Disorder", "Bipolar Disorder", "Schizophrenia", "PTSD", "OCD", "ADHD", "Autism Spectrum Disorder", "Sleep Disorders", "Substance Abuse Disorder"
]

# Expanded 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):
    """Generate expert analysis using Groq API"""
    specialist_title = specialist_mapping[condition][0]
    
    # Construct the prompt
    prompt = f"""As a {specialist_title.lower()}, explain {condition} to a patient with: "{symptoms}".
Structure clearly in 4 parts:
1. Biological process
2. Immediate treatments  
3. Long-term care
4. Emergency signs
Use simple language. Never reference studies/articles."""
    
    # Get response from Groq API
    chat_completion = groq_client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        model="mixtral-8x7b-32768",  # Use Groq's Mixtral model
        temperature=0.5,
        max_tokens=1024,
        top_p=0.9,
        stop=["\n\n\n"]  # Corrected parameter name
    )
    
    # Process response
    raw_analysis = chat_completion.choices[0].message.content
    
    # Clean up output
    clean_text = re.sub(r'\(PMID:\d+\)|\[.*?\]', '', raw_analysis)  # Remove citations
    clean_text = re.sub(r'\b(article|study|volume|issue)\b.*?\.', '', clean_text, flags=re.I)
    sections = [s for s in clean_text.split('\n\n') if any(kw in s for kw in ["process", "treatment", "manage", "emergency"])]
    
    return f"**Expert Analysis:**\n\n" + "\n\n".join(sections[:4])

def create_medical_report(symptoms, language="en"):
    try:
        # Translation
        if language != "en":
            translated = translator(symptoms)[0]['translation_text']
        else:
            translated = symptoms

        # Diagnosis
        result = classifier(translated, conditions, multi_label=False)
        diagnosis = result['labels'][0]
        specialist, system = specialist_mapping.get(diagnosis, ("General Physician", "General Medicine"))
        
        # Generate expert analysis
        expert_analysis = generate_expert_analysis(diagnosis, translated)
        
        # Compile report
        full_report = f"""**Medical Report**
Patient Symptoms: {translated}
Primary Diagnosis: {diagnosis}
Affected System: {system}
Consult: {specialist}

**Key Questions for Your Doctor:**
1. Is this condition acute or chronic?
2. What medication options are suitable for me?
3. What lifestyle changes would help manage this condition?
4. What warning signs require immediate attention?

{expert_analysis}"""
        
        # PDF Generation
        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

# Gradio Interface
interface = gr.Interface(
    fn=create_medical_report,
    inputs=[
        gr.Textbox(label="Describe Your Symptoms", placeholder="e.g., shortness of breath, chest pain..."),
        gr.Dropdown(["en", "hi"], label="Language", value="en")
    ],
    outputs=[
        gr.Textbox(label="Complete Medical Report"),
        gr.File(label="Download PDF Report")
    ],
    title="MedExpert: AI-Powered Clinical Analysis System",
    description="Get comprehensive medical reports with expert-level analysis in any language"
)

interface.launch()