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

# Read HTML content from `re.html`
with open("templates/re.html", "r", encoding="utf-8") as file:
    html_content = file.read()

# 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
with gr.Blocks() as app:
    gr.HTML(html_content)  # Display `re.html` content in Gradio
    gr.Markdown("## Bone Fracture Detection System")
    
    with gr.Row():
        name = gr.Textbox(label="Patient Name")
        age = gr.Number(label="Age")
        gender = gr.Radio(["Male", "Female", "Other"], label="Gender")

    with gr.Row():
        xray1 = gr.Image(type="filepath", label="Upload X-ray Image 1")
        xray2 = gr.Image(type="filepath", label="Upload X-ray Image 2")

    submit_button = gr.Button("Generate Report")
    output_file = gr.File(label="Download Report")

    submit_button.click(
        generate_report,
        inputs=[name, age, gender, xray1, xray2],
        outputs=[output_file],
    )

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