import streamlit as st
import numpy as np
from PIL import Image
import tensorflow as tf
import tensorflow_hub as hub
import tempfile
import urllib.request

# APP user interface configuration
st.set_page_config(page_title="Dog Breed Identifier", layout="centered")

# Add my logo
col1, col2, col3 = st.columns([1, 2, 1])

with col2:
    st.markdown("##") # Add vertical spacing
    st.image("assets/MLOwl_ca_logo_no_bkg_black_cropped.png", width=300)

# Load the model


@st.cache_resource
def load_model():
    url = "https://huggingface.co/turtlemb/dogID_app_model/resolve/main/dog_breed_ID_batch32_cache_prefetch.keras"
    with tempfile.NamedTemporaryFile(suffix=".keras") as tmp:
        urllib.request.urlretrieve(url, tmp.name)
        model = tf.keras.models.load_model(tmp.name, custom_objects={"KerasLayer": hub.KerasLayer})
    return model

model = load_model()

# Define the class names (120 breeds)
class_names = np.array([
    'affenpinscher', 'afghan_hound', 'african_hunting_dog', 'airedale',
    'american_staffordshire_terrier', 'appenzeller', 'australian_terrier',
    'basenji', 'basset', 'beagle', 'bedlington_terrier', 'bernese_mountain_dog',
    'black-and-tan_coonhound', 'blenheim_spaniel', 'bloodhound', 'bluetick',
    'border_collie', 'border_terrier', 'borzoi', 'boston_bull',
    'bouvier_des_flandres', 'boxer', 'brabancon_griffon', 'briard',
    'brittany_spaniel', 'bull_mastiff', 'cairn', 'cardigan',
    'chesapeake_bay_retriever', 'chihuahua', 'chow', 'clumber',
    'cocker_spaniel', 'collie', 'curly-coated_retriever', 'dandie_dinmont',
    'dhole', 'dingo', 'doberman', 'english_foxhound', 'english_setter',
    'english_springer', 'entlebucher', 'eskimo_dog', 'flat-coated_retriever',
    'french_bulldog', 'german_shepherd', 'german_short-haired_pointer',
    'giant_schnauzer', 'golden_retriever', 'gordon_setter', 'great_dane',
    'great_pyrenees', 'greater_swiss_mountain_dog', 'groenendael',
    'ibizan_hound', 'irish_setter', 'irish_terrier', 'irish_water_spaniel',
    'irish_wolfhound', 'italian_greyhound', 'japanese_spaniel', 'keeshond',
    'kelpie', 'kerry_blue_terrier', 'komondor', 'kuvasz',
    'labrador_retriever', 'lakeland_terrier', 'leonberg', 'lhasa',
    'malamute', 'malinois', 'maltese_dog', 'mexican_hairless',
    'miniature_pinscher', 'miniature_poodle', 'miniature_schnauzer',
    'newfoundland', 'norfolk_terrier', 'norwegian_elkhound',
    'norwich_terrier', 'old_english_sheepdog', 'otterhound', 'papillon',
    'pekinese', 'pembroke', 'pomeranian', 'pug', 'redbone',
    'rhodesian_ridgeback', 'rottweiler', 'saint_bernard', 'saluki',
    'samoyed', 'schipperke', 'scotch_terrier', 'scottish_deerhound',
    'sealyham_terrier', 'shetland_sheepdog', 'shih-tzu', 'siberian_husky',
    'silky_terrier', 'soft-coated_wheaten_terrier',
    'staffordshire_bullterrier', 'standard_poodle', 'standard_schnauzer',
    'sussex_spaniel', 'tibetan_mastiff', 'tibetan_terrier', 'toy_poodle',
    'toy_terrier', 'vizsla', 'walker_hound', 'weimaraner',
    'welsh_springer_spaniel', 'west_highland_white_terrier', 'whippet',
    'wire-haired_fox_terrier', 'yorkshire_terrier'
])

# Preprocessing the image


def preprocess(image: Image.Image):
    image = image.resize((224, 224))
    array = np.array(image) / 255.0
    return np.expand_dims(array, axis=0)


# App user interface
with st.container():
    st.title("🐕 The DogID App")
    st.markdown("<p style = 'text-align: left; color: gray;'> by Martin Bijloos | MLOwl.ca</p>",
                unsafe_allow_html=True)
    st.write("Upload a photo of a dog and get the breed prediction!")

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

    if uploaded_file:
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Uploaded Image", use_column_width=True)

        if st.button("Classify"):
            st.info("Processing...")
            input_tensor = preprocess(image)
            prediction = model.predict(input_tensor)[0]

            top_k = 3
            top_indices = prediction.argsort()[-top_k:][::-1]
            top_classes = class_names[top_indices]
            top_confidences = prediction[top_indices]

            st.success("Top Predictions:")
            for breed, score in zip(top_classes, top_confidences):
                st.write(
                    f"- **{breed.replace('_', ' ').title()}**: {score:.2%}")