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

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

# Define image size
image_size = (224, 224)

# Function to detect fracture
def detect_fracture(xray):
    img = Image.open(xray).convert("RGB").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]
    bbox = (50, 50, 150, 150) if prediction > 0.5 else None  # Placeholder
    return prediction, bbox, img

# Function to analyze severity
def analyze_injury(prediction):
    if prediction < 0.3:
        severity = "Mild"
        treatment = "Rest, pain relievers, and follow-up X-ray."
        gov_cost = "₹2,000 - ₹5,000"
        private_cost = "₹10,000 - ₹20,000"
    elif 0.3 <= prediction < 0.7:
        severity = "Moderate"
        treatment = "Plaster cast or splint; possible minor surgery."
        gov_cost = "₹8,000 - ₹15,000"
        private_cost = "₹30,000 - ₹60,000"
    else:
        severity = "Severe"
        treatment = "Major surgery with metal implants, extensive physiotherapy."
        gov_cost = "₹20,000 - ₹50,000"
        private_cost = "₹1,00,000+"

    return severity, treatment, gov_cost, private_cost

# Function to generate report
def generate_report(name, age, gender, xray1, xray2):
    prediction1, bbox1, img1 = detect_fracture(xray1)
    prediction2, bbox2, img2 = detect_fracture(xray2)

    avg_prediction = (prediction1 + prediction2) / 2
    diagnosed_class = "Fractured" if avg_prediction > 0.5 else "Normal"
    severity, treatment, gov_cost, private_cost = analyze_injury(avg_prediction)

    # Draw bounding box if fracture detected
    if bbox1:
        draw = ImageDraw.Draw(img1)
        draw.rectangle(bbox1, outline="red", width=5)

    if bbox2:
        draw = ImageDraw.Draw(img2)
        draw.rectangle(bbox2, outline="red", width=5)

    # Save images
    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", 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, img1_path, img2_path, diagnosed_class, severity, treatment, gov_cost, private_cost

# UI Components
with gr.Blocks() as app:
    gr.Markdown("# 🏥 Bone Fracture Detection & Medical Report")
    gr.Markdown(
        "### A radiologist is a doctor who specializes in reading medical images like X-rays, MRIs, and CT scans to diagnose diseases and injuries."
    )
    gr.Image("x.jpg", label="X-Ray Example")
    gr.Markdown(
        """
        ## **Understanding Bone Fractures**
        - **Closed (Simple):** Bone doesn't pierce skin.
        - **Open (Compound):** Bone breaks skin.
        - **Hairline:** Small stress fracture.
        - **Comminuted:** Bone shatters into pieces.
        - **Avulsion:** Tendon pulls bone fragment.
        - **Compression:** Bones forced together.
        """
    )
    gr.Markdown(
        """
        ## **First Aid**
        - Immobilize the injured area.
        - Control bleeding, cover wounds.
        - Don't straighten broken bones.
        - Use splints, slings for support.
        - Apply cold packs.
        - Seek emergency help.
        """
    )

    with gr.Row():
        gr.Markdown("## 📝 Patient Information Form")

    with gr.Column():
        name = gr.Textbox(label="Patient Name")
        age = gr.Number(label="Age")
        gender = gr.Radio(["Male", "Female", "Other"], label="Gender")

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

    submit_button = gr.Button("Generate Report")

    output_file = gr.File(label="Download Report")
    xray1_output = gr.Image(label="X-ray 1 with Fracture Highlight")
    xray2_output = gr.Image(label="X-ray 2 with Fracture Highlight")
    diagnosis_output = gr.Textbox(label="Fracture Detected")
    severity_output = gr.Textbox(label="Injury Severity")
    treatment_output = gr.Textbox(label="Recommended Treatment")
    gov_cost_output = gr.Textbox(label="Estimated Cost (Govt Hospital)")
    private_cost_output = gr.Textbox(label="Estimated Cost (Private Hospital)")

    submit_button.click(
        generate_report,
        inputs=[name, age, gender, xray1, xray2],
        outputs=[output_file, xray1_output, xray2_output, diagnosis_output, severity_output, treatment_output, gov_cost_output, private_cost_output],
    )

# Run app
if __name__ == "__main__":
    app.launch()