import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import os

# -------------------------
# Load model
# -------------------------
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "cnn_11_layer.h5")
model = tf.keras.models.load_model(MODEL_PATH)

IMG_SIZE = (224, 224)
class_names = ["glioma", "meningioma", "notumor", "pituitary"]

# -------------------------
# Helper: preprocess
# -------------------------
def preprocess(img):
    img = img.resize(IMG_SIZE)
    img = np.array(img) / 255.0

    if img.ndim == 2:
        img = np.stack((img,) * 3, axis=-1)
    if img.shape[-1] == 1:
        img = np.concatenate([img] * 3, axis=-1)

    return np.expand_dims(img, axis=0)

# -------------------------
# Sample images
# -------------------------
SAMPLE_DIR = os.path.join(BASE_DIR, "samples")

sample_options = {
    "None (I'll upload my own)": None,
    "Sample 1": os.path.join(SAMPLE_DIR, "img1.jpg"),
    "Sample 2": os.path.join(SAMPLE_DIR, "img2.jpg"),
    "Sample 3": os.path.join(SAMPLE_DIR, "img3.jpg"),
    "Sample 4": os.path.join(SAMPLE_DIR, "img4.jpg"),
}

# -------------------------
# UI
# -------------------------
st.title("🧠 Brain Tumor Classification")
st.write("Choose a sample image **or upload your own MRI scan**.")

# Select sample
choice = st.selectbox("Choose a sample image:", list(sample_options.keys()))

# File upload
uploaded_file = None
if choice == "None (I'll upload my own)":
    uploaded_file = st.file_uploader("Upload MRI Image...", type=["jpg", "jpeg", "png"])
else:
    uploaded_file = sample_options[choice]

# Display chosen image
if uploaded_file:
    if isinstance(uploaded_file, str):         # Sample path
        image = Image.open(uploaded_file)
    else:                                      # User upload
        image = Image.open(uploaded_file)

    st.image(image, caption="Selected Image", use_column_width=True)

    # Predict button
    if st.button("🔍 Predict Tumor Type"):
        with st.spinner("Analyzing..."):
            img = preprocess(image)
            preds = model.predict(img)
            cls = np.argmax(preds)
            confidence = np.max(preds)

        st.success(f"### Prediction: **{class_names[cls].upper()}**")
        st.info(f"Confidence: **{confidence * 100:.2f}%**")

        st.subheader("Class Probabilities")
        for i, prob in enumerate(preds[0]):
            st.write(f"{class_names[i]}: **{prob*100:.2f}%**")