import streamlit as st
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
from PIL import Image

# 🔧 Configure Streamlit
st.set_page_config(
    page_title="WildVision 🐾 | Animal Identifier",
    layout="centered",
    initial_sidebar_state="auto"
)

# 📦 Load model and class labels
@st.cache_resource
def load_model_once():
    return load_model("animal_model.keras")

@st.cache_data
def load_labels():
    return np.load("class_labels.npy", allow_pickle=True).item()

model = load_model_once()
class_indices = load_labels()
class_labels = list(class_indices.keys())

# 🏷 App Header
st.title("🤖Smart Animal Identifier")
st.markdown("Upload up to **3 animal images** and discover what species they are!")

# 📤 Upload Images
uploaded_files = st.file_uploader(
    "📸 Upload your animal photos",
    type=["jpg", "jpeg", "png"],
    accept_multiple_files=True
)

if uploaded_files:
    # Limit to 3 images
    uploaded_files = uploaded_files[:3]
    cols = st.columns(len(uploaded_files))

    for idx, uploaded_file in enumerate(uploaded_files):
        with cols[idx]:
            st.markdown(f"#### 📷 Image {idx + 1}")

            # 🖼️ Load and preview
            image = Image.open(uploaded_file).convert("RGB")
            preview = image.copy()
            preview.thumbnail((150, 150))
            st.image(preview, caption="Image Preview", use_container_width=True)

            # 🧪 Preprocess for prediction
            resized = image.resize((128, 128))
            img_array = img_to_array(resized) / 255.0
            img_array = np.expand_dims(img_array, axis=0)

            # 🔮 Make prediction
            preds = model.predict(img_array)[0]

            # ✅ Safe check
            if len(preds) != len(class_labels):
                st.error("⚠️ Mismatch between model outputs and label classes.")
            else:
                top_indices = preds.argsort()[-3:][::-1]
                top_labels = [class_labels[i] for i in top_indices]
                top_scores = [preds[i] for i in top_indices]

                st.markdown("#### 🧠 Top Predictions:")
                for label, score in zip(top_labels, top_scores):
                    st.write(f"➡️ **{label.capitalize()}**: {score:.2%}")

                st.success(f"🎯 **Most Likely:** *{top_labels[0].capitalize()}* ")