import os

import streamlit as st

from utils import load_model, load_index
import numpy as np
import matplotlib.pyplot as plt


def app(model_name):
    images_directory = 'images/val2017'
    features_directory = f'features/val2017/{model_name}.tsv'

    files, index = load_index(features_directory)
    model, processor = load_model(f'koclip/{model_name}')

    st.title("Text to Image Search Engine")
    st.markdown("""
        This demonstration explores capability of KoCLIP as a Korean-language Image search engine. Embeddings for each of
        5000 images from [MSCOCO](https://cocodataset.org/#home) 2017 validation set was generated using trained KoCLIP 
        vision model. They are ranked based on cosine similarity distance from input Text query embeddings and top 10 images
        are displayed below.
        
        KoCLIP is a retraining of OpenAI's CLIP model using 82,783 images from [MSCOCO](https://cocodataset.org/#home) dataset and
        Korean caption annotations. Korean translation of caption annotations were obtained from [AI Hub](https://aihub.or.kr/keti_data_board/visual_intelligence).
        Base model `koclip` uses `klue/roberta` as text encoder and `openai/clip-vit-base-patch32` as image encoder.
        Larger model `koclip-large` uses `klue/roberta` as text encoder and bigger `google/vit-large-patch16-224` as image encoder.
        
        Example Queries : 아파트(Apartment), 자동차(Car), 컴퓨터(Computer)
    """)

    query = st.text_input("한글 질문을 적어주세요 (Korean Text Query) :", value="아파트")
    if st.button("질문 (Query)"):
        proc = processor(text=[query], images=None, return_tensors="jax", padding=True)
        vec = np.asarray(model.get_text_features(**proc))
        ids, dists = index.knnQuery(vec, k=10)
        result_files = map(lambda id: files[id], ids)
        result_imgs, result_captions = [], []
        for file, dist in zip(result_files, dists):
            result_imgs.append(plt.imread(os.path.join(images_directory, file)))
            result_captions.append("{:s} (유사도: {:.3f})".format(file, 1.0 - dist))

        st.image(result_imgs[:3], caption=result_captions[:3], width=200)
        st.image(result_imgs[3:6], caption=result_captions[3:6], width=200)
        st.image(result_imgs[6:9], caption=result_captions[6:9], width=200)
        st.image(result_imgs[9:], caption=result_captions[9:], width=200)