import streamlit as st
import numpy as np
import json

from tensorflow.keras.models import Model
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
from tensorflow.keras.preprocessing import image

st.set_page_config(
    page_title="Elephant Species Classifier",
    page_icon="🐘",
    layout="centered"
)

st.title("🐘 Elephant Species Classifier")
st.write("Upload an elephant image and click **Predict**.")

# ------------------ Load Model ------------------ #
@st.cache_resource
def load_artifacts():

    base = MobileNetV2(
        input_shape=(224, 224, 3),
        include_top=False,
        weights="imagenet"
    )

    base.trainable = False

    x = GlobalAveragePooling2D()(base.output)
    x = Dense(448, activation="relu")(x)
    x = Dropout(0.4)(x)
    outputs = Dense(2, activation="softmax")(x)

    model = Model(base.input, outputs)

    model.load_weights("Models/best_Mobilenetv2.weights.h5")

    with open("class_indices.json") as f:
        class_indices = json.load(f)

    idx_to_class = {v: k for k, v in class_indices.items()}

    return model, idx_to_class


model, idx_to_class = load_artifacts()

# ------------------ Upload ------------------ #
uploaded_file = st.file_uploader(
    "Choose an Elephant Image",
    type=["jpg", "jpeg", "png"]
)

if uploaded_file is not None:

    col1, col2, col3 = st.columns([1,2,1])

    with col2:
        st.image(uploaded_file, width=250, caption="Uploaded Image")

    st.write("")

    if st.button("🔍 Predict", type="primary", use_container_width=True):

        with st.spinner("Predicting..."):

            img = image.load_img(uploaded_file, target_size=(224,224))
            x = image.img_to_array(img)
            x = x / 255.0
            x = np.expand_dims(x, axis=0)

            preds = model.predict(x, verbose=0)

            pred_idx = np.argmax(preds)
            pred_class = idx_to_class[pred_idx]
            confidence = np.max(preds) * 100

        st.success(f"### 🐘 Prediction: {pred_class}")
        st.info(f"**Confidence:** {confidence:.2f}%")

else:
    st.info("Please upload an image to begin.")