import numpy as np
import streamlit as st
from PIL import Image

from models.deep_learning import DeepLearningGradCAM, ImageEmbeddings
from utils import configs
from utils.functional import (
    generate_empty_space,
    get_default_images,
    get_most_salient_object,
    set_page_config,
    set_seed,
)

# Set seed
set_seed()

# Set page config
set_page_config("Image Embeddings with Deep Learning", "📈")

# Sidebar
name_model = st.sidebar.selectbox("Select Model", tuple(configs.NAME_MODELS.keys()))
support_set_method = st.sidebar.selectbox(
    "Select Support Set Method", configs.SUPPORT_SET_METHODS
)
freeze_model = st.sidebar.checkbox("Freeze Model", value=True)
pretrained_model = st.sidebar.checkbox("Pretrained Model", value=True)


# Load Model
@st.cache_resource
def load_model(
    name_model: str, support_set_method: str, freeze_model: bool, pretrained_model: bool
):
    image_embeddings = ImageEmbeddings(
        name_model, freeze_model, pretrained_model, support_set_method
    )
    custom_grad_cam = DeepLearningGradCAM(
        name_model, freeze_model, pretrained_model, support_set_method
    )
    return image_embeddings, custom_grad_cam


image_embeddings, custom_grad_cam = load_model(
    name_model, support_set_method, freeze_model, pretrained_model
)

# Application Description
st.markdown("# ❓ Application Description")
st.write(
    """
    Introducing our Image Embeddings with Deep Learning app! 📈 This app is designed to help you generate meaningful representations of your images using the power of deep learning. With just a few clicks, you can easily extract high-dimensional feature vectors from your images, which can be used for a variety of tasks such as image retrieval, clustering, and classification.

Our user-friendly interface makes it easy for anyone to get started. Simply upload your images and let our deep learning models do the rest. You can visualize the resulting embeddings, explore their similarities and differences, and even use them for downstream applications like clustering or classification.

Whether you're a data scientist, a machine learning enthusiast, or just curious about the power of deep learning, our Image Embeddings with Deep Learning app is sure to be a valuable addition to your toolkit. So why wait? Give it a try today and start exploring the rich world of image embeddings!
    """
)

uploaded_file = st.file_uploader(
    "Upload image file ", type=["jpg", "jpeg", "png", "bmp", "tiff"]
)
select_default_images = st.selectbox("Select default images ", get_default_images())
st.caption("Default Images will be used if no image is uploaded.")
select_image_button = st.button("Select Image ")
if select_image_button:
    st.success("Image selected")

if select_image_button and uploaded_file is not None:
    image = np.array(Image.open(uploaded_file).convert("RGB"))
    st.session_state["image"] = image
elif select_image_button and uploaded_file is None:
    image = np.array(Image.open(select_default_images).convert("RGB"))
    st.session_state["image"] = image

if st.session_state.get("image") is not None:
    image = st.session_state.get("image")
    col, col2, col3 = st.columns(3)
    col2.write("## 📸 Preview Image ")
    col2.image(image, use_column_width=True)
    predict_image_button = col2.button("Get Image Embeddings")
    generate_empty_space(2)
    if predict_image_button:
        with st.spinner("Getting Image Embeddings..."):
            result_embeddings = image_embeddings.get_embeddings(image)
            result_grad_cam = custom_grad_cam.get_grad_cam(image)
            inference_time = result_embeddings["inference_time"]
            col1, col2, col3 = st.columns(3)
            col1.write("### 🙂 Source Image")
            col1.image(image, use_column_width=True)
            col2.write("### 😎 Grad CAM Image")
            col2.image(result_grad_cam, use_column_width=True)
            col3.write("### 🤔 Most Salient Object")
            col3.image(get_most_salient_object(image), use_column_width=True)
            st.write("### 📝 Result")
            st.json(
                {
                    "image_embeddings": result_embeddings["embeddings"].tolist(),
                }
            )
            st.write(f"Image Embeddings Shape: {result_embeddings['embeddings'].shape}")
            st.write(f"Inference Time: {inference_time:.2f} s")
            st.session_state["image"] = None