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
from reportlab.lib import colors
from reportlab.platypus import Table, TableStyle

# 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, allergies, cause, 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_details = {
        "Mild": "Your fracture is classified as **Mild**. It may heal with rest, pain relievers, and a follow-up X-ray. Avoid excessive movement of the affected area.",
        "Moderate": "Your fracture is classified as **Moderate**. You may require a plaster cast, splint, or minor surgery. Recovery takes **4-8 weeks**.",
        "Severe": "Your fracture is classified as **Severe**. Surgery with metal implants and extensive physiotherapy is required. Recovery takes **several months** with proper rehabilitation."
    }
    treatment = treatment_details[severity]

    # Estimated cost & duration
    cost_duration_data = [
        ["Hospital Type", "Estimated Cost", "Recovery Time"],
        ["Government Hospital", f"₹{2000 if severity == 'Mild' else 8000 if severity == 'Moderate' else 20000} - ₹{5000 if severity == 'Mild' else 15000 if severity == 'Moderate' else 50000}", "4-12 weeks"],
        ["Private Hospital", f"₹{10000 if severity == 'Mild' else 30000 if severity == 'Moderate' else 100000}+", "6 weeks - Several months"]
    ]

    # Save X-ray images for report
    img1 = Image.open(xray1).resize((300, 300))
    img2 = Image.open(xray2).resize((300, 300))
    img1_path = f"{name}_xray1.png"
    img2_path = f"{name}_xray2.png"
    img1.save(img1_path)
    img2.save(img2_path)

    # Generate PDF report
    report_path = f"{name}_fracture_report.pdf"
    c = canvas.Canvas(report_path, pagesize=letter)
    c.setFont("Helvetica-Bold", 14)
    c.drawString(200, 770, "Bone Fracture Detection Report")

    # Patient details
    c.setFont("Helvetica", 12)
    c.drawString(100, 740, f"Patient Name: {name}")
    c.drawString(100, 720, f"Age: {age}")
    c.drawString(100, 700, f"Gender: {gender}")
    c.drawString(100, 680, f"Allergies: {allergies if allergies else 'None'}")
    c.drawString(100, 660, f"Cause of Injury: {cause if cause else 'Not Provided'}")

    # Diagnosis
    c.setFont("Helvetica-Bold", 12)
    c.drawString(100, 630, "Diagnosis & Treatment Plan:")
    c.setFont("Helvetica", 11)
    c.drawString(100, 610, f"Fracture Detected: {diagnosed_class}")
    c.drawString(100, 590, f"Injury Severity: {severity}")
    c.setFont("Helvetica", 10)
    c.drawString(100, 570, f"{treatment}")

    # Load and insert X-ray images
    c.drawInlineImage(img1_path, 50, 250, width=250, height=250)
    c.drawInlineImage(img2_path, 320, 250, width=250, height=250)

    # Cost estimation table
    table = Table(cost_duration_data)
    table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
        ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
        ('GRID', (0, 0), (-1, -1), 1, colors.black)
    ]))
    table.wrapOn(c, 400, 300)
    table.drawOn(c, 100, 150)

    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():
        allergies = gr.Textbox(label="Allergies (if any)")
        cause = gr.Textbox(label="Cause of Injury")

    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, allergies, cause, xray1, xray2],
        outputs=[output_file],
    )

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