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, UnidentifiedImageError

# Configure Streamlit
st.set_page_config(
    page_title="Animal Classifier",
    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())

# Title
st.title("Animal Image Classifier")
st.markdown("Upload up to **3 images** to get predictions")

# File uploader
uploaded_files = st.file_uploader(
    "Upload animal images",
    type=["jpg", "jpeg", "png"],
    accept_multiple_files=True
)

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

    for idx, uploaded_file in enumerate(uploaded_files):
        with cols[idx]:
            st.markdown(f"### Image {idx + 1}")
            try:
                image = Image.open(uploaded_file).convert("RGB")
                preview = image.copy()
                preview.thumbnail((150, 150))
                st.image(preview, caption="Preview", use_container_width=False)

                # Resize and preprocess
                resized = image.resize((128, 128))
                img_array = img_to_array(resized) / 255.0
                img_array = np.expand_dims(img_array, axis=0)

                # Prediction
                preds = model.predict(img_array)[0]
                top_index = np.argmax(preds)
                top_label = class_labels[top_index]

                # Display top prediction only
                st.success(f"Top Prediction: **{top_label.capitalize()}**")

            except UnidentifiedImageError:
                st.error("Error: Invalid image file.")
            except Exception as e:
                st.error(f"Unexpected error: {str(e)}")