import streamlit as st
import random
import os
import clip
from funcs.get_similarity import get_similarity_score, create_filelist, load_embeddings, find_matches
from funcs.fiass_similaruty import load_embeddings, encode_text, find_matches_fiass
import torch.nn.functional as F
import pandas as pd
import zipfile
from unzip import unzip

zip_filename = 'img.zip'
target_dir = 'photos'

# Create the target directory if it doesn't exist
if not os.path.exists(target_dir):
    os.makedirs(target_dir)

# Open the zip file and extract its contents to the target directory
with zipfile(zip_filename, 'r') as zip_ref:
    zip_ref.extractall(target_dir)

device = 'cpu'
model_path = "weights/ViT-B-32.pt"

model, preprocess = clip.load('ViT-B/32', device)

file_name = create_filelist('img')
features = load_embeddings('embeddings/emb_images_5000.npy')
df = pd.read_csv('data/results.csv')

random_queries = ['friends playing cards', 'rock band playing on guitars', 'policeman cross the road', 
                  'sleeping kids', 'football team playing on the grass' , 'learning programming'
                  ]

st.header('Find my pic!')

request = st.text_input('Write a description of the picture', ' Two people at the photo')

if request:
    unzip('img.zip')

img_count = st.slider('How much pic you need?', 4, 8, 6, 2)


matches = find_matches_fiass(features, request, file_name, img_count)
row1, row2 = st.columns(2)
    
if st.button('Find!'):

    selected_filenames = matches

    for i in range(int(img_count/2)):
        filename = selected_filenames[i]
        img_path = filename
        img_discription = df[df['image_name'] == filename.split('/')[1]]['comment'].iloc[0]
        with row1:
            st.image(img_path, width=300, caption=img_discription)

    # display next 3 images in the second row
    for i in range(int(img_count/2), img_count):
        filename = selected_filenames[i]
        img_path = filename
        img_discription = df[df['image_name'] == filename.split('/')[1]]['comment'].iloc[0]
        with row2:
            st.image(img_path, width=300, caption=img_discription)