import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force TensorFlow to use CPU

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

# Load the trained model
model = tf.keras.models.load_model("my_keras_model.h5")

# Function to process X-rays and generate a PDF report
def generate_report(name, age, gender, xray1, xray2):
    image_size = (224, 224)

    def predict_fracture(xray_path):
        img = Image.open(xray_path).resize(image_size)
        img_array = image.img_to_array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        prediction = model.predict(img_array)[0][0]
        return prediction

    # Predict on both X-rays
    prediction1 = predict_fracture(xray1)
    prediction2 = predict_fracture(xray2)
    avg_prediction = (prediction1 + prediction2) / 2
    diagnosed_class = "Fractured" if avg_prediction > 0.5 else "Normal"

    # Injury severity classification
    severity = "Mild" if avg_prediction < 0.3 else "Moderate" if avg_prediction < 0.7 else "Severe"
    treatment = {
        "Mild": "Rest, pain relievers, follow-up X-ray.",
        "Moderate": "Plaster cast, possible minor surgery.",
        "Severe": "Major surgery, metal implants, physiotherapy."
    }[severity]
    gov_cost = {"Mild": "₹2,000 - ₹5,000", "Moderate": "₹8,000 - ₹15,000", "Severe": "₹20,000 - ₹50,000"}[severity]
    private_cost = {"Mild": "₹10,000 - ₹20,000", "Moderate": "₹30,000 - ₹60,000", "Severe": "₹1,00,000+"}[severity]

    # Generate PDF report
    report_path = f"{name}_fracture_report.pdf"
    c = canvas.Canvas(report_path, pagesize=letter)
    c.setFont("Helvetica", 12)
    c.drawString(100, 750, f"Patient Name: {name}")
    c.drawString(100, 730, f"Age: {age}")
    c.drawString(100, 710, f"Gender: {gender}")
    c.drawString(100, 690, f"Diagnosis: {diagnosed_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()

    return report_path  # Return path for auto-download

# 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."
)

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