import json
import numpy as np
import streamlit as st
from st_clickable_images import clickable_images

from clip_multilingual.search import MultiLingualSearch
from clip_multilingual.models import Tokenizer


@st.cache(
    suppress_st_warning=True,
    hash_funcs={
        Tokenizer: lambda _: None
    }
)
def load_model():
    unsplash_base_folder = './'
    all_embeddings = np.load(f'{unsplash_base_folder}/embeddings.npy')
    with open(f'{unsplash_base_folder}/urls.json') as f:
        all_urls = json.load(f)
    return MultiLingualSearch(all_embeddings, all_urls)

semantic_search = load_model()

description = '''
# Multilingual Semantic Search

**Search images in any language powered by OpenAI's [CLIP](https://openai.com/blog/clip/) and XMLRoberta.**
'''

st.sidebar.markdown(description)
examples = [
     'pessoas sorrindo',
     '微笑的人',
     'насмејани људи',
]
for example in examples:
    if st.sidebar.button(example):
        st.session_state.query = example

_, c, _ = st.columns((1, 3, 1))
if "query" in st.session_state:
    query = c.text_input("", value=st.session_state["query"])
else:
    query = c.text_input("", value="people smiling")
if len(query) > 0:
    results = semantic_search.search(query)
    clicked = clickable_images(
        [result['image'] for result in results],
        titles=[f'Prob: {result["prob"]}' for result in results],
        div_style={
            "display": "flex",
            "justify-content": "center",
            "flex-wrap": "wrap",
        },
        img_style={"margin": "2px", "height": "200px"},
    )