import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing import image
from PIL import Image
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import os

# Load the trained model once
model = tf.keras.models.load_model("my_keras_model.h5")
image_size = (224, 224)  # Ensure consistent image size

# Function to analyze injury severity
def analyze_injury(prediction):
    if prediction < 0.3:
        return "Mild", "Rest and pain relief.", "₹2,000 - ₹5,000", "₹10,000 - ₹20,000"
    elif 0.3 <= prediction < 0.7:
        return "Moderate", "Plaster cast or minor surgery.", "₹8,000 - ₹15,000", "₹30,000 - ₹60,000"
    else:
        return "Severe", "Major surgery with metal implants.", "₹20,000 - ₹50,000", "₹1,00,000+"

# Function to generate report
def generate_report(patient_name, age, gender, xray1_path, xray2_path):
    if not os.path.exists(xray1_path) or not os.path.exists(xray2_path):
        return "Error: One or both X-ray images are missing!"

    try:
        # Process X-ray 1
        img1 = Image.open(xray1_path).resize(image_size).convert("RGB")
        img_array1 = image.img_to_array(img1)
        img_array1 = np.expand_dims(img_array1, axis=0) / 255.0  
        prediction1 = model.predict(img_array1)[0][0]

        # Process X-ray 2
        img2 = Image.open(xray2_path).resize(image_size).convert("RGB")
        img_array2 = image.img_to_array(img2)
        img_array2 = np.expand_dims(img_array2, axis=0) / 255.0  
        prediction2 = model.predict(img_array2)[0][0]

        # Get final analysis
        avg_prediction = (prediction1 + prediction2) / 2
        predicted_class = "Fractured" if avg_prediction > 0.5 else "Normal"
        severity, treatment, gov_cost, private_cost = analyze_injury(avg_prediction)

        # Generate PDF Report
        report_path = f"{patient_name}_fracture_report.pdf"
        c = canvas.Canvas(report_path, pagesize=letter)
        c.setFont("Helvetica", 12)
        c.drawString(100, 750, f"Patient Name: {patient_name}")
        c.drawString(100, 730, f"Age: {age}")
        c.drawString(100, 710, f"Gender: {gender}")
        c.drawString(100, 690, f"Diagnosis: {predicted_class}")
        c.drawString(100, 670, f"Injury Severity: {severity}")
        c.drawString(100, 650, f"Recommended Treatment: {treatment}")
        c.drawString(100, 630, f"Estimated Cost (Govt Hospital): {gov_cost}")
        c.drawString(100, 610, f"Estimated Cost (Private Hospital): {private_cost}")
        c.save()

        if os.path.exists(report_path):
            return report_path
        else:
            return "Error: Report generation failed!"

    except Exception as e:
        return f"Error generating report: {str(e)}"

# Define Gradio Interface
interface = gr.Interface(
    fn=generate_report,
    inputs=[
        gr.Textbox(label="Patient Name"),
        gr.Number(label="Age"),
        gr.Radio(["Male", "Female", "Other"], label="Gender"),
        gr.Image(type="filepath", label="Upload X-ray Image 1"),
        gr.Image(type="filepath", label="Upload X-ray Image 2"),
    ],
    outputs=gr.File(label="Download Report"),
    title="Bone Fracture Detection & Medical Report",
    description="Enter patient details, upload two X-ray images, and generate a detailed medical report with treatment suggestions and cost estimates."
)

if __name__ == "__main__":
    interface