import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# ---------------- LOAD TRAINED MODEL ---------------- #
model = tf.keras.models.load_model("chest_xray_model.h5")  # Make sure this file is uploaded
class_labels = ["Normal", "Pneumonia"]

# ---------------- PREDICTION FUNCTION ---------------- #
def predict_xray(img):
    # Preprocess
    img = img.resize((224, 224))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    
    # Model prediction
    prediction = model.predict(img_array, verbose=0)[0][0]
    label = class_labels[int(prediction > 0.5)]
    confidence = prediction if prediction > 0.5 else 1 - prediction
    
    # Detailed preliminary radiology report and first aid
    if label == "Pneumonia":
        report = (
            "Preliminary Radiology Report:\n"
            "- The chest X-ray shows opacities or infiltrates consistent with pneumonia.\n"
            "- Findings suggest possible lung inflammation or infection.\n"
            "- Further diagnostic tests (blood tests, sputum culture, oxygen saturation) recommended.\n\n"
            "First Aid / Immediate Actions:\n"
            "1. Seek medical attention immediately for confirmation and treatment.\n"
            "2. Monitor for severe symptoms: high fever, shortness of breath, chest pain, confusion.\n"
            "3. Ensure hydration and rest.\n"
            "4. Avoid self-medicating with antibiotics without doctor supervision.\n"
            "5. Use a mask and maintain good hygiene to prevent spread if contagious.\n"
            "6. Keep a pulse oximeter if available; seek emergency care if oxygen saturation < 94%.\n"
            "7. Note any worsening symptoms and report them to healthcare providers promptly.\n"
        )
    else:
        report = (
            "Preliminary Radiology Report:\n"
            "- No visible signs of pneumonia detected on this X-ray.\n"
            "- Lungs appear clear, but clinical correlation is advised.\n\n"
            "General Advice:\n"
            "1. Maintain healthy habits: good hydration, nutrition, and regular exercise.\n"
            "2. Seek medical attention if respiratory symptoms develop.\n"
            "3. Continue monitoring for cough, fever, or shortness of breath.\n"
        )
    
    return f"Prediction: {label} ({confidence*100:.2f}% confidence)\n\n{report}"

# ---------------- GRADIO INTERFACE ---------------- #
interface = gr.Interface(
    fn=predict_xray,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Pneumonia Detector, Analyser and What to do",
    description="Upload a chest X-ray to get a detailed preliminary report and first-aid recommendations."
)

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