import os
import gradio as gr
from fpdf import FPDF
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders
import torch
from torchvision import transforms
from PIL import Image

# Set environment variable to disable GPU if needed
os.environ["CUDA_VISIBLE_DEVICES"] = ""

# Load the trained fracture detection model
model = torch.load("my_keras_model.h5")  
model.eval()

# Function to predict fracture
def predict_fracture(xray):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  
        transforms.ToTensor(),
    ])
    image = transform(xray).unsqueeze(0)
    with torch.no_grad():
        output = model(image)
        predicted_class = "Fractured" if torch.argmax(output) == 1 else "Not Fractured"
        confidence = torch.nn.functional.softmax(output, dim=1).max().item() * 100
    return predicted_class, confidence

# Function to generate PDF report
def generate_report(name, age, gender, weight, height, allergies, injury_cause, address, parent_name, email, xray):
    # Ensure input limits
    name = name[:50] if name else "N/A"
    age = str(age) if age else "N/A"
    gender = gender if gender else "N/A"
    weight = str(weight) + " kg" if weight else "N/A"
    height = str(height) + " cm" if height else "N/A"
    allergies = allergies[:100] if allergies else "None"
    injury_cause = injury_cause[:500] if injury_cause else "Not specified"
    address = address[:150] if address else "N/A"
    parent_name = parent_name[:50] if parent_name else "N/A"
    
    # Fake hospital details
    hospital_name = "CityCare Orthopedic Hospital"
    hospital_address = "123 Medical Lane, Health City, Country"

    # Predict fracture
    prediction, confidence = predict_fracture(xray)

    # Create PDF
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    
    # Title
    pdf.set_font("Arial", style="B", size=14)
    pdf.cell(200, 10, hospital_name, ln=True, align="C")
    pdf.set_font("Arial", size=10)
    pdf.cell(200, 5, hospital_address, ln=True, align="C")
    pdf.ln(10)

    # Patient Information
    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(200, 10, "Patient Report", ln=True, align="C")
    pdf.ln(5)

    pdf.set_font("Arial", size=10)
    pdf.cell(200, 5, f"Patient Name: {name}", ln=True)
    pdf.cell(200, 5, f"Age: {age}  |  Gender: {gender}", ln=True)
    pdf.cell(200, 5, f"Weight: {weight}  |  Height: {height}", ln=True)
    pdf.cell(200, 5, f"Allergies: {allergies}", ln=True)
    pdf.cell(200, 5, f"Cause of Injury: {injury_cause}", ln=True)
    pdf.cell(200, 5, f"Address: {address}", ln=True)
    pdf.cell(200, 5, f"Parent/Guardian: {parent_name}", ln=True)
    pdf.ln(10)

    # X-ray image
    if xray:
        pdf.set_font("Arial", style="B", size=12)
        pdf.cell(200, 10, "X-ray Image", ln=True, align="C")
        pdf.ln(5)
        xray_path = "temp_xray.png"
        xray.save(xray_path)
        pdf.image(xray_path, x=40, w=130)
        os.remove(xray_path)
        pdf.ln(5)

        # Prediction result
        pdf.set_font("Arial", style="B", size=10)
        pdf.cell(200, 5, f"Prediction: {prediction} (Confidence: {confidence:.2f}%)", ln=True, align="C")

    pdf.ln(10)

    # Diagnosis and Recommendation
    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(200, 10, "Diagnosis & Recommendations", ln=True)
    pdf.set_font("Arial", size=10)
    pdf.multi_cell(0, 5, "Based on the provided X-ray and details, the following suggestions are recommended:")
    
    pdf.set_font("Arial", style="I", size=10)
    pdf.cell(200, 5, "- Immediate medical consultation is advised.", ln=True)
    pdf.cell(200, 5, "- Pain management with prescribed medications.", ln=True)
    pdf.cell(200, 5, "- Possible surgical intervention if required.", ln=True)
    pdf.cell(200, 5, "- Rest and immobilization of the affected area.", ln=True)
    pdf.cell(200, 5, "- Follow-up X-ray and rehabilitation therapy.", ln=True)
    
    pdf.ln(5)
    pdf.set_font("Arial", style="B", size=10)
    pdf.cell(200, 5, "Estimated Treatment Costs:", ln=True)
    
    pdf.set_font("Arial", size=10)
    pdf.cell(200, 5, "Government Hospital: $500 - $1,200", ln=True)
    pdf.cell(200, 5, "Private Hospital: $2,000 - $5,000", ln=True)

    # Save PDF
    pdf_path = "patient_report.pdf"
    pdf.output(pdf_path)

    # Send email
    send_email(email, name, hospital_name, pdf_path)

    return pdf_path

# Function to send email with PDF report
def send_email(email, patient_name, hospital_name, pdf_path):
    sender_email = "your_email@gmail.com"
    sender_password = "your_app_password"  # Use App Password

    subject = f"Patient Report - {patient_name}"
    
    message = MIMEMultipart()
    message["From"] = sender_email
    message["To"] = email
    message["Subject"] = subject
    body = f"Dear {patient_name},\n\nYour medical report from {hospital_name} is attached. Please review the details and consult a doctor if needed.\n\nBest regards,\n{hospital_name}"
    message.attach(MIMEText(body, "plain"))

    with open(pdf_path, "rb") as attachment:
        part = MIMEBase("application", "octet-stream")
        part.set_payload(attachment.read())
        encoders.encode_base64(part)
        part.add_header("Content-Disposition", f"attachment; filename={pdf_path}")
        message.attach(part)

    try:
        server = smtplib.SMTP("smtp.gmail.com", 587)
        server.starttls()
        server.login(sender_email, sender_password)
        server.sendmail(sender_email, email, message.as_string())
        server.quit()
        print("Email sent successfully!")
    except Exception as e:
        print(f"Error sending email: {e}")

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# Bone Fracture Detection & Diagnosis")
    gr.Markdown("Upload an X-ray, enter patient details, and get a report with treatment suggestions.")

    xray = gr.Image(label="Upload X-ray", type="pil", value="samples/sample_xray.jpg")

    submit = gr.Button("Generate Report")
    output = gr.File()

    submit.click(generate_report, [name, age, gender, weight, height, allergies, injury_cause, address, parent_name, email, xray], output)

app.launch()