import streamlit as st
st.set_page_config(page_title="PinoyPaws", layout="wide")

import tensorflow as tf
from PIL import Image
import numpy as np
import os
import json
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.utils import plot_model
import io

# === Page Configuration ===
st.sidebar.title("📍 PinoyPaws Navigation")

# === Sidebar Navigation ===
page = st.sidebar.selectbox("Navigate to", ["Overview", "Predict Breed", "Model Insights"])

with st.sidebar.expander("ℹ️ About this App"):
    st.markdown("Built with 🐍 TensorFlow and 🧠 MobileNetV2")

# === Load model ===
@st.cache_resource
def load_model():
    model_path = os.path.join("src", "model", "dog_breed_classifier.h5")
    model = tf.keras.models.load_model(model_path)
    return model

model = load_model()

# === Load class names ===
@st.cache_data
def load_class_names():
    labels_path = os.path.join("src", "model", "class_names.json")
    with open(labels_path, "r") as f:
        return json.load(f)

class_names = load_class_names()

# === Preprocess image ===
def preprocess_image(image: Image.Image) -> np.ndarray:
    image = image.resize((224, 224))
    image_array = np.array(image)
    if image_array.shape[-1] == 4:
        image_array = image_array[..., :3]
    image_array = preprocess_input(image_array)
    return np.expand_dims(image_array, axis=0)

# === Page: Overview ===
if page == "Overview":
    st.title("🐾 PinoyPaws: Dog Breed Classifier")
    st.markdown("""
    Welcome to **PinoyPaws**, a dog breed classifier tailored to recognize common dog breeds found in the Philippines 🐕🇵🇭.

    ### 📌 Features:
    - 📷 Upload a dog image and let our AI guess the breed!
    - 🧠 Built using **MobileNetV2** for fast and lightweight inference
    - 📊 Confidence score included
    - 🐕 Trained on 5 local and common breeds:
                - **Beagle**
                - **Chihuahua**
                - **Golden Retriever**
                - **Shih Tzu**
                - **Siberian Husky**

    ### 📁 Input:
    - Accepts `.jpg`, `.jpeg`, `.png` images
    - Optimized for images where the dog is clearly visible

    You can get started by choosing **Predict Breed** in the sidebar.
    """)

# === Page: Predict Breed ===
elif page == "Predict Breed":
    st.title("🔎🐶 Predict Dog Breed")
    st.write(f"Upload an image of a dog and let the model predict its breed from {len(class_names)} common dog breeds.")

    uploaded_file = st.file_uploader("📷 Choose a dog image...", type=["jpg", "jpeg", "png"])

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

            if st.button("Predict"):

                with st.spinner("Classifying..."):
                    try:

                        input_tensor = preprocess_image(image)
                        prediction = model.predict(input_tensor)
                        predicted_index = int(np.argmax(prediction))
                        predicted_class = class_names[predicted_index]
                        confidence = np.max(prediction)

                        st.success(f"🐶 Predicted Breed: **{predicted_class}**")
                        st.info(f"📊 Confidence: {confidence * 100:.2f}%")

                    except Exception as e:
                        st.error(f"An error occurred: {e}")

# === Page: Model Insights ===
elif page == "Model Insights":
    st.title("📊 Model Insights & Architecture")

    st.markdown("### 🧠 Model Summary")
    string_io = io.StringIO()
    model.summary(print_fn=lambda x: string_io.write(x + "\n"))
    summary_str = string_io.getvalue()
    st.text(summary_str)

    st.markdown("### 🧬 Model Details")
    st.write(f"• Total parameters: `{model.count_params():,}`")
    st.write("• Architecture: **MobileNetV2** base with custom dense layers")

    st.markdown("### 📚 Classes Detected")
    st.write(f"The model can classify the following {len(class_names)} breeds:")
    st.markdown(" - " + "\n - ".join(class_names))