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

from models.deep_learning import DeepLearningGradCAM, ImageSimilarity
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 Similarity 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_similarity = ImageSimilarity(
        name_model, freeze_model, pretrained_model, support_set_method
    )
    custom_grad_cam = DeepLearningGradCAM(
        name_model, freeze_model, pretrained_model, support_set_method
    )
    return image_similarity, custom_grad_cam


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

# Application Description
st.markdown("# ❓ Application Description")
st.write(
    """
    Looking for a powerful and efficient way to find similar images? Look no further than our Image Similarity with Deep Learning application! Using the latest in deep learning technology, this app allows you to quickly and easily identify images that are similar in content, style, and more.

Whether you're a professional photographer, a graphic designer, or just someone who loves taking pictures, our app can help you find the images you need in a flash. With a user-friendly interface and advanced algorithms, it's never been easier to search through large collections of images and find the ones that best match your needs.

So why wait? Give our Image Similarity with Deep Learning app a try today and start exploring the world of images like never before! 🖼️
    """
)

col1, col2 = st.columns(2)
uploaded_file1 = col1.file_uploader(
    "Upload image file 1", type=["jpg", "jpeg", "png", "bmp", "tiff"]
)
select_default_images1 = col1.selectbox("Select default images 1", get_default_images())
col1.caption("Default Images 1 will be used if no image is uploaded.")
select_image_button1 = col1.button("Select Image 1")
if select_image_button1:
    st.success("Image 1 selected")

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

if select_image_button1 and uploaded_file1 is not None:
    image1 = np.array(Image.open(uploaded_file1).convert("RGB"))
    st.session_state["image1"] = image1
elif select_image_button1 and uploaded_file1 is None:
    image1 = np.array(Image.open(select_default_images1).convert("RGB"))
    st.session_state["image1"] = image1

if select_image_button2 and uploaded_file2 is not None:
    image2 = np.array(Image.open(uploaded_file2).convert("RGB"))
    st.session_state["image2"] = image2
elif select_image_button2 and uploaded_file2 is None:
    image2 = np.array(Image.open(select_default_images2).convert("RGB"))
    st.session_state["image2"] = image2

if (
    st.session_state.get("image1") is not None
    and st.session_state.get("image2") is not None
):
    image1 = st.session_state.get("image1")
    image2 = st.session_state.get("image2")
    col1, col2 = st.columns(2)
    col1.write("## 📸 Preview Image 1")
    col1.image(image1, use_column_width=True)
    col2.write("## 📸 Preview Image 2")
    col2.image(image2, use_column_width=True)
    predict_image_button = st.button("Get Image Similarity")
    generate_empty_space(2)
    if predict_image_button:
        with st.spinner("Getting Image Similarity..."):
            result_similarity = image_similarity.get_similarity(image1, image2)
            result_grad_cam1 = custom_grad_cam.get_grad_cam(image1)
            result_grad_cam2 = custom_grad_cam.get_grad_cam(image2)
            inference_time = result_similarity["inference_time"]
            col1, col2 = st.columns(2)
            col1.write("### 😎 Grad CAM Image 1")
            col1.image(result_grad_cam1, use_column_width=True)
            col2.write("### 😎 Grad CAM Image 2")
            col2.image(result_grad_cam2, use_column_width=True)
            col1, col2 = st.columns(2)
            col1.write("### 🤔 Most Salient Object Image 1")
            col1.image(get_most_salient_object(image1), use_column_width=True)
            col2.write("### 🤔 Most Salient Object Image 2")
            col2.image(get_most_salient_object(image2), use_column_width=True)
            st.write("### 📝 Result")
            st.write(f"Similarity Score: {result_similarity['similarity'] * 100:.2f}%")
            st.write(f"Similarity Label: {result_similarity['result_similarity'].title()}")
            st.write(f"Inference Time: {inference_time:.2f} s")
            st.session_state["image1"] = None
            st.session_state["image2"] = None