import streamlit as st
import torch
import os

from PIL import Image
from utils.search import search_images

device = "cuda" if torch.cuda.is_available() else "cpu"
st.set_page_config(layout="wide")

search_path = None

def show_image(result):
    col1, col2, col3 = st.columns(3)

    image_folder2 = "./data/images_mr"
    for idx, image_name in enumerate(result.ids):
        if idx % 3 == 0:
            with col1:
                file_name = str(image_name) + ".jpg"
                image_path = os.path.join(image_folder2, file_name)
                st.image(image_path, caption=image_name, width=200)
        elif idx % 3 == 1:
            with col2:
                file_name = str(image_name) + ".jpg"
                image_path = os.path.join(image_folder2, file_name)
                st.image(image_path, caption=image_name, width=200)
        else:
            with col3:
                file_name = str(image_name) + ".jpg"
                image_path = os.path.join(image_folder2, file_name)
                st.image(image_path, caption=image_name, width=200)
                
image_folder1 = "./examples"
image_paths = []
for file_name in os.listdir(image_folder1):
    image_paths.append(os.path.join(image_folder1, file_name))
# st.write(image_paths)

with st.sidebar:
    if st.sidebar.button("Choose a examples"):
        search_path = image_paths[0]
        st.image(image_paths[0], caption="example", width=150)
    
    search_term = st.file_uploader(label="Chose a file", type=["jpg", "png"])
    if search_term is None:
        st.text("Please upload a image!")
    else:
        image = Image.open(search_term).convert('RGB')
        st.image(image, width=300)
    if search_term:
        button = st.sidebar.button("Search")
        if button:
            search_path = search_term


st.header("Image Retrieval")
st.write("This is a simple app for image retrieval using Resnet18 and Vector Database")



if search_path is not None:
    result = search_images(search_path)
    show_image(result)