import streamlit as st
import numpy as np
from PIL import Image

import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.resnet50 import preprocess_input

# -----------------------------------------------------------
# 1. Config
# -----------------------------------------------------------

IMG_SIZE = (128, 128)  # zelfde als bij training
MODEL_PATH = "indian_birds_resnet50.h5"   # <-- pas aan indien nodig

# VUL HIER JOUW LABELS IN (dezelfde volgorde als tijdens training)
labels = [
    "Asian Green Bee-Eater",
    "Brown-Headed Barbet",
    "Cattle Egret",
    "Common Kingfisher",
    "Common Myna",
    "Common Rosefinch",
    "Common Tailorbird",
    "Coppersmith Barbet",
    "Forest Wagtail",
    "Gray Wagtail",
    "Hoopoe",
    "House Crow",
    "Indian Grey Hornbill",
    "Indian Peafowl",
    "Indian Pitta",
    "Indian Roller",
    "Jungle Babbler",
    "Northern Lapwing",
    "Red Wattled Lapwing",
    "Ruddy Shelduck",
    "Rufous Treepie",
    "Sarus Crane",
    "White Wagtail",
    "White-Breasted Kingfisher",
    "White-Breasted Waterhen"
]

label_map = {name: i for i, name in enumerate(labels)}
inv_label_map = {v: k for k, v in label_map.items()}

# -----------------------------------------------------------
# 2. Model laden (gecached)
# -----------------------------------------------------------

@st.cache_resource
def load_bird_model():
    model = load_model(MODEL_PATH)
    return model

model = load_bird_model()

# -----------------------------------------------------------
# 3. Hulpfuncties
# -----------------------------------------------------------

def preprocess_image(img: Image.Image) -> np.ndarray:
    """Resize + preprocess zoals bij training."""
    img = img.convert("RGB").resize(IMG_SIZE)
    x = np.array(img).astype("float32")
    x = preprocess_input(x)         # ResNet50 preprocessing
    x = np.expand_dims(x, axis=0)   # batch-dimensie
    return x

def predict_image(img: Image.Image):
    x = preprocess_image(img)
    probs = model.predict(x, verbose=0)[0]      # vorm (25,)
    pred_id = int(np.argmax(probs))
    pred_label = inv_label_map[pred_id]
    pred_prob = float(probs[pred_id])

    # top-3
    top3_ids = probs.argsort()[-3:][::-1]
    top3 = [(inv_label_map[int(i)], float(probs[i])) for i in top3_ids]

    return pred_label, pred_prob, top3

# -----------------------------------------------------------
# 4. Streamlit UI
# -----------------------------------------------------------

st.set_page_config(page_title="Indian Bird Classifier", layout="centered")

st.title("🕊️ Indian Bird Species Classifier")
st.write(
    "Upload een vogelafbeelding uit de *25 Indian Bird Species* dataset "
    "en het model (ResNet50) voorspelt de soort."
)

uploaded_file = st.file_uploader(
    "Kies een afbeelding (.jpg, .png)",
    type=["jpg", "jpeg", "png"]
)

if uploaded_file is not None:
    # Afbeelding tonen
    image = Image.open(uploaded_file)

    st.image(image, caption="Geüploade afbeelding", use_column_width=True)

    if st.button("🔮 Voorspel vogelsoort"):
        with st.spinner("Model is aan het voorspellen..."):
            pred_label, pred_prob, top3 = predict_image(image)

        st.success(f"Voorspelling: **{pred_label}** ({pred_prob:.2%} zekerheid)")

        st.subheader("Top 3 voorspellingen")
        for name, p in top3:
            st.write(f"- {name}: **{p:.2%}**")

else:
    st.info("Upload een afbeelding om te starten.")