import gradio as gr
from tensorflow import keras
import numpy as np
from PIL import Image
from fpdf import FPDF  # Import library to generate PDF

# Load the pre-trained Keras model
model = keras.models.load_model('retino_model.keras')  

# Define class names  for your model predictions
class_names = ['Healthy', 'Mild DR', 'Moderate DR', 'Proliferative DR', 'Severe DR']

# Function to provide additional care information based on the predicted condition
def eye_care_recommendations(predicted_class):
    recommendations = {
        'Healthy': 'Your eyes seem healthy. Continue with regular eye check-ups and maintain a balanced diet.',
        'Mild DR': 'Mild signs of diabetic retinopathy. Ensure strict blood sugar control and regular eye exams.',
        'Moderate DR': 'Moderate diabetic retinopathy detected. Consult with an ophthalmologist for treatment options.',
        'Proliferative DR': 'Advanced stage detected. Immediate medical attention is required to prevent further vision loss.',
        'Severe DR': 'Severe diabetic retinopathy detected. Medical intervention is necessary. Please visit a doctor immediately.'
    }
    return recommendations.get(predicted_class, "No recommendation available.")

# Function to generate a PDF report
def generate_pdf(patient_info, prediction, care_info):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, txt="Diabetic Retinopathy Prediction Report", ln=True, align='C')
    
    pdf.ln(10)
    pdf.cell(200, 10, txt=f"Patient Name: {patient_info['name']}", ln=True)
    pdf.cell(200, 10, txt=f"Age: {patient_info['age']}", ln=True)
    
    pdf.ln(10)
    pdf.cell(200, 10, txt=f"Prediction: {prediction}", ln=True)
    pdf.cell(200, 10, txt=f"Eye Care Recommendations: {care_info}", ln=True)
    
    # Save PDF
    pdf_output = f"{patient_info['name']}_report.pdf"
    pdf.output(pdf_output)
    
    return pdf_output

# Prediction function that processes the image and returns the result and care advice
def predict(image, name, age):
    # Resize image to the expected size for the model
    image = image.resize((128, 128))  # Adjust this size based on your model's input size
    image = np.expand_dims(np.array(image), axis=0)  # Add batch dimension

    # Make a prediction
    predictions = model.predict(image)
    predicted_class_index = np.argmax(predictions, axis=1)[0]
    predicted_class = class_names[predicted_class_index]

    # Get eye care recommendations based on the prediction
    care_info = eye_care_recommendations(predicted_class)

    # Generate a PDF report
    patient_info = {'name': name, 'age': age}
    pdf_report = generate_pdf(patient_info, predicted_class, care_info)

    # Return the prediction, care information, and PDF link
    return f"Predicted Condition: {predicted_class}", care_info, pdf_report

# Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Patient Name"), gr.Textbox(label="Age")],
    outputs=[gr.Textbox(label="Prediction"), gr.Textbox(label="Eye Care Recommendations"), gr.File(label="Download Report (PDF)")],
    title="Diabetic Retinopathy Prediction",
    description="Upload a retinal image, enter patient information, and the model will predict the stage of Diabetic Retinopathy. Eye care recommendations and a PDF report will be provided."
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()