import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
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 and get bounding box
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]
    
    # Dummy bounding box for now (assuming model enhancement needed)
    # Ideally, the model should return bounding box coordinates
    bbox = (50, 50, 150, 150) if prediction > 0.5 else None  # Placeholder
    
    return prediction, bbox, img

# Function to analyze injury 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):
    # Analyze X-ray images
    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 is 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 modified 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

# 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="file", label="Upload X-ray Image 1"),
        gr.Image(type="file", label="Upload X-ray Image 2"),
    ],
    outputs=[
        gr.File(label="Download Report"),
        gr.Image(label="X-ray 1 with Fracture Highlight"),
        gr.Image(label="X-ray 2 with Fracture Highlight"),
        gr.Textbox(label="Fracture Detected"),
        gr.Textbox(label="Injury Severity"),
        gr.Textbox(label="Recommended Treatment"),
        gr.Textbox(label="Estimated Cost (Govt Hospital)"),
        gr.Textbox(label="Estimated Cost (Private Hospital)")
    ],
    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 app
if __name__ == "__main__":
    interface.launch()