import streamlit as st
import numpy as np
import tensorflow as tf
import pydicom
import cv2

# --------------------------------------------------
# App Config
# --------------------------------------------------
st.set_page_config(page_title="Pneumonia Detection", layout="centered")

st.title("🫁 Pneumonia Detection from Chest X-ray")
st.write("Upload a chest X-ray (DICOM) to detect pneumonia.")

# --------------------------------------------------
# Load Serialized Model (Inference Only)
# --------------------------------------------------
@st.cache_resource
def load_model():
    return tf.keras.models.load_model(
        "pneumonia_final_model.keras",
        compile=False   # avoids optimizer warnings
    )

model = load_model()

# --------------------------------------------------
# Preprocess DICOM
# --------------------------------------------------
def preprocess_dicom(uploaded_file):
    dicom = pydicom.dcmread(uploaded_file)
    img = dicom.pixel_array.astype(np.float32)

    img = (img - img.min()) / (img.max() - img.min() + 1e-6)
    img = cv2.resize(img, (224, 224))

    # grayscale → RGB
    img = np.stack([img] * 3, axis=-1)
    img = np.expand_dims(img, axis=0)

    return img

# --------------------------------------------------
# File Upload
# --------------------------------------------------
uploaded_file = st.file_uploader(
    "Upload Chest X-ray (.dcm file)",
    type=["dcm"]
)

# --------------------------------------------------
# Prediction
# --------------------------------------------------
if uploaded_file is not None:
    try:
        image = preprocess_dicom(uploaded_file)

        with st.spinner("Analyzing X-ray..."):
            probs = model.predict(image)[0]

        # ---- 3-class → binary mapping ----
        pneumonia_prob = float(probs[2])
        non_pneumonia_prob = float(probs[0] + probs[1])

        if pneumonia_prob >= 0.5:
            prediction = "Pneumonia"
            confidence = pneumonia_prob
            st.error(f"⚠️ {prediction} detected")
        else:
            prediction = "Non-Pneumonia"
            confidence = non_pneumonia_prob
            st.success(f"✅ {prediction}")

        st.subheader("Prediction Result")
        st.write(f"**Prediction:** {prediction}")
        st.write(f"**Confidence:** {confidence:.2%}")

        # Optional: show raw probabilities (for transparency)
        with st.expander("View raw class probabilities"):
            st.write({
                "Normal": float(probs[0]),
                "No Lung Opacity": float(probs[1]),
                "Lung Opacity (Pneumonia)": float(probs[2])
            })

    except Exception as e:
        st.error(f"Error processing file: {e}")