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

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

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

from tensorflow import keras
from tensorflow.keras.layers import TFSMLayer

# Load the model as a layer (in the SavedModel format)
#eye_model = TFSMLayer('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",
    },
    "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",
    },
    "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",
    },
    "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"

# Loading the custom model for consultation with the doctor
try:
    # Force using the slow tokenizer for compatibility
    tokenizer = AutoTokenizer.from_pretrained("ahmed-7124/NeuraMedAW", use_fast=False)
except Exception as e:
    print(f"Tokenizer error: {e}")
    tokenizer = AutoTokenizer.from_pretrained("ahmed-7124/NeuraMedAW", use_fast=False)

model = AutoModelForCausalLM.from_pretrained("ahmed-7124/NeuraMedAW")

def consult_doctor(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=100)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Functions for the app

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_or_extract_report(patient_id, pdf=None, report_text=None):
    if pdf:
        # Extract text from PDF
        with pdfplumber.open(pdf.name) as pdf_file:
            report_text = "".join([page.extract_text() for page in pdf_file.pages])
    
    if not report_text:
        return "❌ Please provide a report text or upload a PDF."

    # Analyze the report
    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}\n💊 Medication: {medication}\n⚠ Precaution: {precaution}\n👩‍⚕ Recommended Doctor: {doctor}"


# 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(input_image)
#     labels = ['Cataract', 'Conjunctivitis', 'Glaucoma', 'Normal']
#     confidence_scores = {labels[i]: round(predictions[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",
# )

from transformers import AutoTokenizer, AutoModelForCausalLM

try:
    tokenizer = AutoTokenizer.from_pretrained("ahmed-7124/dgptAW")
    model = AutoModelForCausalLM.from_pretrained("ahmed-7124/dgptAW")
    print("Model and tokenizer loaded successfully!")
except Exception as e:
    print(f"Error loading tokenizer or model: {e}")
    print("Trying with GPT-2 tokenizer as a fallback...")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    model = AutoModelForCausalLM.from_pretrained("ahmed-7124/dgptAW")

from transformers import AutoTokenizer, AutoModelForCausalLM

def answer_medical_query(query):
    try:
        inputs = tokenizer(query, return_tensors="pt")
        outputs = model.generate(
            inputs.input_ids,
            max_length=50,
            temperature=0.7,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id  # Avoid errors if padding is required
        )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    except Exception as e:
        return f"An error occurred while generating a response: {e}"



# Unified Gradio Interface
analyze_report_interface = gr.Interface(
    fn=analyze_or_extract_report,
    inputs=[
        gr.Number(label="Patient ID"),
        gr.File(label="Upload PDF Report"),  # Removed optional=True
        gr.Textbox(label="Report Text (Optional)"),
    ],
    outputs="text",
)

# eye_disease_interface = gr.Interface(
#     fn=predict_eye_disease,
#     inputs=gr.Image(label="Upload an Eye Image", type="numpy"),
#     outputs="text",
# )

# Gradio Interface for Medical Query Answering
medical_query_interface = gr.Interface(
    fn=answer_medical_query,
    inputs=gr.Textbox(label="Ask a Medical Term"),
    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",
)
# Gradio Interface for Medical Query Answering
medical_query_interface = gr.Interface(
    fn=answer_medical_query,
    inputs=gr.Textbox(label="Ask a Medical Term"),
    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",
)

consult_doctor_interface = gr.Interface(
    fn=consult_doctor,
    inputs=gr.Textbox(label="Enter Your Query for the Doctor"),
    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("Analyze Medical Report"):
        analyze_report_interface.render()

    with gr.Tab("Doctor Consult"):
        consult_doctor_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("Encyclopedia"):
        medical_query_interface.render()

    with gr.Tab("Patient Dashboard"):
        patient_dashboard_interface.render()
    
    with gr.Tab("Doctor Dashboard"):
        doctor_dashboard_interface.render()
    

# Launch the app
app.launch(share=True)