import gradio as gr
import tensorflow as tf
import pdfplumber
from transformers import pipeline
import timm
import torch
import pandas as pd

# Load pre-trained zero-shot model for text classification
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Pre-trained ResNet50 model for X-ray or image analysis
image_model = timm.create_model('resnet50', pretrained=True)
image_model.eval()

# Load saved TensorFlow eye disease detection model
#eye_model = tf.keras.models.load_model('model.h5')

# Patient database
patients_db = []

# Disease details for medical report analyzer
disease_details = {
    "anemia": {
        "medication": (
            "Iron supplements (e.g., ferrous sulfate), "
            "Vitamin B12 injections (for pernicious anemia), "
            "Folic acid supplements."
        ),
        "precaution": (
            "Consume iron-rich foods like spinach, red meat, and lentils. "
            "Pair iron-rich foods with vitamin C to enhance absorption. "
            "Avoid tea or coffee with meals as they inhibit iron absorption."
        ),
        "doctor": "Hematologist",
    },
    "viral infection": {
        "medication": (
            "Antiviral drugs (e.g., oseltamivir for flu, acyclovir for herpes). "
            "Over-the-counter medications for symptom relief, such as ibuprofen for fever and body aches."
        ),
        "precaution": (
            "Stay hydrated by drinking plenty of fluids. "
            "Isolate to prevent spreading the infection. "
            "Rest adequately to support recovery. Maintain proper hygiene."
        ),
        "doctor": "Infectious Disease Specialist",
    },
    "liver disease": {
        "medication": (
            "Hepatoprotective drugs (e.g., ursodeoxycholic acid, silymarin). "
            "Antiviral therapy for viral hepatitis. "
            "Diuretics for managing fluid retention (e.g., spironolactone)."
        ),
        "precaution": (
            "Avoid alcohol and hepatotoxic drugs. "
            "Follow a low-fat diet and avoid processed foods. "
            "Regularly monitor liver function tests."
        ),
        "doctor": "Hepatologist",
    },
    "diabetes": {
        "medication": (
            "Oral hypoglycemics (e.g., metformin). "
            "Insulin therapy for Type 1 diabetes or advanced Type 2 diabetes. "
            "GLP-1 receptor agonists (e.g., liraglutide) for improving blood sugar control."
        ),
        "precaution": (
            "Monitor blood glucose levels daily. "
            "Follow a low-carb, high-fiber diet. "
            "Engage in regular physical activity. "
            "Avoid sugary foods and beverages."
        ),
        "doctor": "Endocrinologist",
    },
    "hypertension": {
        "medication": (
            "ACE inhibitors (e.g., lisinopril). "
            "Beta-blockers (e.g., metoprolol). "
            "Calcium channel blockers (e.g., amlodipine). "
            "Diuretics (e.g., hydrochlorothiazide)."
        ),
        "precaution": (
            "Reduce salt intake to less than 2g per day. "
            "Engage in at least 150 minutes of moderate exercise weekly. "
            "Avoid smoking and excessive alcohol consumption. "
            "Manage stress through relaxation techniques like yoga or meditation."
        ),
        "doctor": "Cardiologist",
    },
    "pneumonia": {
        "medication": (
            "Antibiotics (e.g., amoxicillin or azithromycin for bacterial pneumonia). "
            "Antiviral therapy if caused by viruses like influenza. "
            "Supplemental oxygen in severe cases."
        ),
        "precaution": (
            "Get plenty of rest and stay hydrated. "
            "Use a humidifier to ease breathing. "
            "Avoid smoking or exposure to pollutants. "
            "Ensure vaccination against influenza and pneumococcus."
        ),
        "doctor": "Pulmonologist",
    },
    "asthma": {
        "medication": (
            "Short-acting bronchodilators (e.g., albuterol) for quick relief. "
            "Inhaled corticosteroids (e.g., fluticasone) for long-term control. "
            "Leukotriene receptor antagonists (e.g., montelukast) for reducing inflammation."
        ),
        "precaution": (
            "Avoid known allergens like pollen, dust, and pet dander. "
            "Carry a rescue inhaler at all times. "
            "Practice breathing exercises to strengthen lungs. "
            "Avoid cold air or strenuous exercise without a warm-up."
        ),
        "doctor": "Pulmonologist",
    },
    "kidney disease": {
        "medication": (
            "ACE inhibitors or ARBs (e.g., losartan) for controlling blood pressure. "
            "Erythropoietin-stimulating agents for anemia management. "
            "Phosphate binders (e.g., sevelamer) to manage high phosphate levels."
        ),
        "precaution": (
            "Limit salt, potassium, and phosphorus in the diet. "
            "Stay hydrated but avoid overhydration. "
            "Avoid NSAIDs and other nephrotoxic drugs. "
            "Monitor kidney function and blood pressure regularly."
        ),
        "doctor": "Nephrologist",
    },
    "thyroid disorder": {
        "medication": (
            "Levothyroxine for hypothyroidism. "
            "Antithyroid medications (e.g., methimazole) for hyperthyroidism. "
            "Beta-blockers for symptomatic relief in hyperthyroidism."
        ),
        "precaution": (
            "Ensure regular thyroid function tests. "
            "Avoid foods that interfere with thyroid hormone absorption (e.g., soy, certain vegetables). "
            "Follow medication schedules precisely without skipping doses."
        ),
        "doctor": "Endocrinologist",
    },
    "arthritis": {
        "medication": (
            "Nonsteroidal anti-inflammatory drugs (NSAIDs) for pain relief. "
            "Disease-modifying antirheumatic drugs (DMARDs) for rheumatoid arthritis. "
            "Biologics (e.g., adalimumab) in severe cases."
        ),
        "precaution": (
            "Engage in low-impact exercises like swimming or yoga. "
            "Use ergonomic furniture to reduce joint strain. "
            "Maintain a healthy weight to reduce joint stress. "
            "Apply heat or cold therapy for symptom relief."
        ),
        "doctor": "Rheumatologist",
    },
    "depression": {
        "medication": (
            "Selective serotonin reuptake inhibitors (SSRIs, e.g., sertraline). "
            "Serotonin-norepinephrine reuptake inhibitors (SNRIs, e.g., venlafaxine). "
            "Tricyclic antidepressants (e.g., amitriptyline) in specific cases."
        ),
        "precaution": (
            "Engage in regular physical exercise. "
            "Maintain a routine and avoid isolation. "
            "Consider therapy (e.g., CBT or psychotherapy). "
            "Avoid alcohol and recreational drugs."
        ),
        "doctor": "Psychiatrist",
    },
}


# Passwords
doctor_password = "doctor123"

# Functions
def register_patient(name, age, gender, password):
    patient_id = len(patients_db) + 1
    patients_db.append({
        "ID": patient_id,
        "Name": name,
        "Age": age,
        "Gender": gender,
        "Password": password,
        "Diagnosis": "",
        "Medications": "",
        "Precautions": "",
        "Doctor": ""
    })
    return f"✅ Patient {name} registered successfully. Patient ID: {patient_id}"

def analyze_report(patient_id, report_text):
    candidate_labels = list(disease_details.keys())
    result = classifier(report_text, candidate_labels)
    diagnosis = result['labels'][0]

    # Update patient's record
    medication = disease_details[diagnosis]['medication']
    precaution = disease_details[diagnosis]['precaution']
    doctor = disease_details[diagnosis]['doctor']
    for patient in patients_db:
        if patient['ID'] == patient_id:
            patient.update(Diagnosis=diagnosis, Medications=medication, Precautions=precaution, Doctor=doctor)
    return f"🔍 Diagnosis: {diagnosis}"

def extract_pdf_report(pdf):
    text = ""
    with pdfplumber.open(pdf.name) as pdf_file:
        for page in pdf_file.pages:
            text += page.extract_text()
    return text

'''def predict_eye_disease(input_image):
    input_image = tf.image.resize(input_image, [224, 224]) / 255.0
    input_image = tf.expand_dims(input_image, 0)
    predictions = eye_model.predict(input_image)
    labels = ['Cataract', 'Conjunctivitis', 'Glaucoma', 'Normal']
    confidence_scores = {labels[i]: round(predictions[0][i] * 100, 2) for i in range(len(labels))}
    if confidence_scores['Normal'] > 50:
        return f"Congrats! No disease detected. Confidence: {confidence_scores['Normal']}%"
    return "\n".join([f"{label}: {confidence}%" for label, confidence in confidence_scores.items()])
'''
def doctor_space(patient_id):
    for patient in patients_db:
        if patient["ID"] == patient_id:
            return f"⚠ Precautions: {patient['Precautions']}\n👩‍⚕ Recommended Doctor: {patient['Doctor']}"
    return "❌ Patient not found. Please check the ID."

def pharmacist_space(patient_id):
    for patient in patients_db:
        if patient["ID"] == patient_id:
            return f"💊 Medications: {patient['Medications']}"
    return "❌ Patient not found. Please check the ID."

def patient_dashboard(patient_id, password):
    for patient in patients_db:
        if patient["ID"] == patient_id and patient["Password"] == password:
            return (f"🩺 Name: {patient['Name']}\n"
                    f"📋 Diagnosis: {patient['Diagnosis']}\n"
                    f"💊 Medications: {patient['Medications']}\n"
                    f"⚠ Precautions: {patient['Precautions']}\n"
                    f"👩‍⚕ Recommended Doctor: {patient['Doctor']}")
    return "❌ Access Denied: Invalid ID or Password."

def doctor_dashboard(password):
    if password != doctor_password:
        return "❌ Access Denied: Incorrect Password"
    if not patients_db:
        return "No patient records available."
    details = []
    for patient in patients_db:
        details.append(f"🩺 Name: {patient['Name']}\n"
                       f"📋 Diagnosis: {patient['Diagnosis']}\n"
                       f"💊 Medications: {patient['Medications']}\n"
                       f"⚠ Precautions: {patient['Precautions']}\n"
                       f"👩‍⚕ Recommended Doctor: {patient['Doctor']}")
    return "\n\n".join(details)

# Gradio Interfaces
registration_interface = gr.Interface(
    fn=register_patient,
    inputs=[
        gr.Textbox(label="Patient Name"),
        gr.Number(label="Age"),
        gr.Radio(label="Gender", choices=["Male", "Female", "Other"]),
        gr.Textbox(label="Set Password", type="password"),
    ],
    outputs="text",
)

pdf_extraction_interface = gr.Interface(
    fn=extract_pdf_report,
    inputs=gr.File(label="Upload PDF Report"),
    outputs="text",
)

report_analysis_interface = gr.Interface(
    fn=analyze_report,
    inputs=[
        gr.Number(label="Patient ID"),
        gr.Textbox(label="Report Text"),
    ],
    outputs="text",
)

'''eye_disease_interface = gr.Interface(
    fn=predict_eye_disease,
    inputs=gr.Image(label="Upload an Eye Image", type="numpy"),
    outputs="text",
)
'''
doctor_space_interface = gr.Interface(
    fn=doctor_space,
    inputs=gr.Number(label="Patient ID"),
    outputs="text",
)

pharmacist_space_interface = gr.Interface(
    fn=pharmacist_space,
    inputs=gr.Number(label="Patient ID"),
    outputs="text",
)

patient_dashboard_interface = gr.Interface(
    fn=patient_dashboard,
    inputs=[
        gr.Number(label="Patient ID"),
        gr.Textbox(label="Password", type="password"),
    ],
    outputs="text",
)

doctor_dashboard_interface = gr.Interface(
    fn=doctor_dashboard,
    inputs=gr.Textbox(label="Doctor Password", type="password"),
    outputs="text",
)

# Gradio App Layout
with gr.Blocks() as app:
    gr.Markdown("# Medico GPT")
    
    with gr.Tab("Patient Registration"):
        registration_interface.render()
    
    with gr.Tab("Analyze Medical Report"):
        report_analysis_interface.render()
    
    with gr.Tab("Extract PDF Report"):
        pdf_extraction_interface.render()
     
    with gr.Tab("Doctor Space"):
        doctor_space_interface.render()
    
    with gr.Tab("Pharmacist Space"):
        pharmacist_space_interface.render()
    
    with gr.Tab("Patient Dashboard"):
        patient_dashboard_interface.render()
    
    with gr.Tab("Doctor Dashboard"):
        doctor_dashboard_interface.render()

app.launch(share=True)