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
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 (TensorFlow Model without Keras)
eye_model = tf.saved_model.load('model')

# Patient database
patients_db = []

# Disease details for medical report analyzer
disease_details = {
    "anemia": {"medication": "Iron supplements", "precaution": "Eat iron-rich foods", "doctor": "Hematologist"},
    "viral infection": {"medication": "Antiviral drugs", "precaution": "Stay hydrated", "doctor": "Infectious Disease Specialist"},
    "liver disease": {"medication": "Hepatoprotective drugs", "precaution": "Avoid alcohol", "doctor": "Hepatologist"},
    "diabetes": {"medication": "Metformin or insulin", "precaution": "Monitor sugar levels", "doctor": "Endocrinologist"},
}

# Passwords
doctor_password = "doctor123"

from transformers import AutoTokenizer, AutoModelForCausalLM

try:
    # Force using the slow tokenizer
    tokenizer = AutoTokenizer.from_pretrained("harishussain12/PastelMed", use_fast=False)
except Exception as e:
    print(f"Tokenizer error: {e}")
    tokenizer = AutoTokenizer.from_pretrained("harishussain12/PastelMed", use_fast=False)

model = AutoModelForCausalLM.from_pretrained("harishussain12/PastelMed")

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
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(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",
)

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",
)

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("Extract PDF Report"):
        pdf_extraction_interface.render()
    
    with gr.Tab("Ophthalmologist Space"):
        eye_disease_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()

    with gr.Tab("Doctor Consult"):
        consult_doctor_interface.render()

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