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

# ================== SETTINGS ==================
IMG_SIZE = (170, 170)  # same as target_size in your ImageDataGenerator

CLASS_NAMES = [
    "Black Rot",
    "ESCA",
    "Healthy",
    "Leaf Blight",
]

MODEL_FILENAME = "grape_disease_model.h5"  # make sure this file is in the same folder
# =============================================


@st.cache_resource
def load_model():
    """
    Load the Keras model once and cache it.
    """
    if not os.path.exists(MODEL_FILENAME):
        st.error(f"Model file '{MODEL_FILENAME}' not found in this directory.")
        st.stop()

    model = tf.keras.models.load_model(MODEL_FILENAME)
    return model


model = load_model()


def preprocess_image(img: Image.Image) -> np.ndarray:
    """
    - Resize to IMG_SIZE
    - Convert to numpy array
    - Normalize (same as rescale=1./255 in your ImageDataGenerator)
    - Add batch dimension
    """
    img = img.resize(IMG_SIZE)
    img_array = np.array(img)

    # same normalization as during training
    img_array = img_array.astype("float32") / 255.0

    # (H, W, 3) -> (1, H, W, 3)
    img_array = np.expand_dims(img_array, axis=0)
    return img_array


def predict_image(img: Image.Image):
    """
    Returns (predicted_label, confidence, all_probabilities)
    for multi-class classification.
    """
    x = preprocess_image(img)
    preds = model.predict(x)

    # expected: shape (1, 4) with softmax
    probs = preds[0]
    idx = int(np.argmax(probs))
    label = CLASS_NAMES[idx]
    confidence = float(probs[idx])
    return label, confidence, probs


# ================ STREAMLIT UI =================
st.set_page_config(
    page_title="Grape Disease Classifier",
    layout="centered"
)

st.title("🍇 Grape Disease Classifier")
st.write(
    "Upload a photo of a grape leaf. "
    "The model predicts whether the leaf is **Black Rot**, **ESCA**, **Healthy**, "
    "or **Leaf Blight**."
)

uploaded_file = st.file_uploader(
    "Upload an image (.jpg, .jpeg or .png)",
    type=["jpg", "jpeg", "png"]
)

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded image", use_column_width=True)

    if st.button("Predict"):
        with st.spinner("Running prediction..."):
            label, confidence, probs = predict_image(image)

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

        st.write("Class probabilities:")
        for name, p in zip(CLASS_NAMES, probs):
            st.write(f"- {name}: {float(p):.2%}")
else:
    st.info("Upload an image above to get a prediction.")