import os
import gradio as gr
import tensorflow as tf
import numpy as np
import pandas as pd
from transformers import pipeline
import pdfplumber
from PIL import Image
import timm
import torch

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

# Pre-trained model for X-ray 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 325mg)",
        "precaution": "Increase intake of iron-rich foods like spinach, red meat, and beans.",
        "doctor": "Hematologist"
    },
    "viral infection": {
        "medication": "Antiviral drugs (e.g., Oseltamivir 75mg for flu)",
        "precaution": "Rest, stay hydrated, avoid close contact with others, and wash hands frequently.",
        "doctor": "Infectious Disease Specialist"
    },
    "liver disease": {
        "medication": "Hepatoprotective drugs (e.g., Ursodeoxycholic acid 300mg)",
        "precaution": "Avoid alcohol and maintain a balanced diet, avoid fatty foods.",
        "doctor": "Hepatologist"
    },
    "kidney disease": {
        "medication": "Angiotensin-converting enzyme inhibitors (e.g., Lisinopril 10mg)",
        "precaution": "Monitor salt intake, stay hydrated, and avoid NSAIDs.",
        "doctor": "Nephrologist"
    },
    "diabetes": {
        "medication": "Metformin (e.g., 500mg) or insulin therapy",
        "precaution": "Follow a low-sugar diet, monitor blood sugar levels, and exercise regularly.",
        "doctor": "Endocrinologist"
    },
    "hypertension": {
        "medication": "Antihypertensive drugs (e.g., Amlodipine 5mg)",
        "precaution": "Reduce salt intake, manage stress, and avoid smoking.",
        "doctor": "Cardiologist"
    },
    "COVID-19": {
        "medication": "Supportive care, antiviral drugs (e.g., Remdesivir 200mg in severe cases)",
        "precaution": "Follow isolation protocols, wear a mask, stay hydrated, and rest.",
        "doctor": "Infectious Disease Specialist"
    },
    "pneumonia": {
        "medication": "Antibiotics (e.g., Amoxicillin 500mg or Azithromycin 250mg)",
        "precaution": "Rest, avoid smoking, stay hydrated, and get proper ventilation.",
        "doctor": "Pulmonologist"
    }
}

# Functions
def register_patient(name, age, gender):
    patient_id = len(patients_db) + 1
    patients_db.append({
        "ID": patient_id,
        "Name": name,
        "Age": age,
        "Gender": gender,
        "Diagnosis": "",
        "Medications": "",
        "Precautions": ""
    })
    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']
    for patient in patients_db:
        if patient['ID'] == patient_id:
            patient.update(Diagnosis=diagnosis, Medications=medication, Precautions=precaution)
    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:
            diagnosis = patient["Diagnosis"]
            medication = patient["Medications"]
            precaution = patient["Precautions"]
            doctor = disease_details.get(diagnosis, {}).get("doctor", "No doctor available")
            return (f"🩺 Patient Name: {patient['Name']}\n"
                    f"📋 Diagnosis: {diagnosis}\n"
                    f"💊 Medications: {medication}\n"
                    f"⚠️ Precautions: {precaution}\n"
                    f"👩‍⚕️ Recommended Doctor: {doctor}")
    return "Patient not found. Please check the ID."

def pharmacist_space(patient_id):
    for patient in patients_db:
        if patient["ID"] == patient_id:
            diagnosis = patient["Diagnosis"]
            medication = patient["Medications"]
            return f"💊 Patient Name: {patient['Name']}\n📋 Prescribed Medications: {medication}"
    return "Patient not found. Please check the ID."

# 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"])], outputs="text")
report_analysis_interface = gr.Interface(fn=analyze_report, inputs=[gr.Number(label="Patient ID"), gr.Textbox(label="Report Text")], outputs="text")
pdf_report_extraction_interface = gr.Interface(fn=extract_pdf_report, inputs=gr.File(label="Upload PDF Report"), outputs="text")
eye_disease_interface = gr.Interface(fn=predict_eye_disease, inputs=gr.Image(label="Upload an Eye Image", type="numpy"), outputs="text")
dashboard_interface = gr.Interface(fn=lambda: pd.DataFrame(patients_db), inputs=None, outputs="dataframe")
doctor_interface = gr.Interface(fn=doctor_space, inputs=gr.Number(label="Patient ID"), outputs="text")
pharmacist_interface = gr.Interface(fn=pharmacist_space, inputs=gr.Number(label="Patient ID"), outputs="text")

# Gradio App Layout
with gr.Blocks() as app:
    gr.Markdown("# Medical Analyzer and Eye Disease Detection")
    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_report_extraction_interface.render()
    with gr.Tab("Detect Eye Disease"):
        eye_disease_interface.render()
    with gr.Tab("Doctor Space"):
        doctor_interface.render()
    with gr.Tab("Pharmacist Space"):
        pharmacist_interface.render()
    with gr.Tab("Patient Dashboard"):
        dashboard_interface.render()

app.launch(share=True)