import pandas as pd
import numpy as np
import clip
import gradio as gr
from utils import *
import os

# Load the open CLIP model
model, preprocess = clip.load("ViT-B/32", device=device)
from pathlib import Path

# Download from Github Releases
if not Path('unsplash-dataset/photo_ids.csv').exists():
    os.system('''wget https://github.com/haltakov/natural-language-image-search/releases/download/1.0.0/photo_ids.csv -O unsplash-dataset/photo_ids.csv''')

if not Path('unsplash-dataset/features.npy').exists():
    os.system('''wget https://github.com/haltakov/natural-language-image-search/releases/download/1.0.0/features.npy - O unsplash-dataset/features.npy''')


# Load the photo IDs
photo_ids = pd.read_csv("unsplash-dataset/photo_ids.csv")
photo_ids = list(photo_ids['photo_id'])

# Load the features vectors
photo_features = np.load("unsplash-dataset/features.npy")

# Convert features to Tensors: Float32 on CPU and Float16 on GPU
if device == "cpu":
  photo_features = torch.from_numpy(photo_features).float().to(device)
else:
  photo_features = torch.from_numpy(photo_features).to(device)

# Print some statistics
print(f"Photos loaded: {len(photo_ids)}")


def search_by_text_and_photo(query_text, query_img, query_photo_id=None, photo_weight=0.5):
    # Encode the search query
    if not query_text and not query_photo_id:
        return []

    text_features = encode_search_query(model, query_text)

    if query_photo_id:
        # Find the feature vector for the specified photo ID
        query_photo_index = photo_ids.index(query_photo_id)
        query_photo_features = photo_features[query_photo_index]

        # Combine the test and photo queries and normalize again
        search_features = text_features + query_photo_features * photo_weight
        search_features /= search_features.norm(dim=-1, keepdim=True)

        # Find the best match
        best_photo_ids = find_best_matches(search_features, photo_features, photo_ids, 10)

    elif query_img:
        query_photo_features = model.encode_image(query_img)
        query_photo_features = query_photo_features / query_photo_features.norm(dim=1, keepdim=True)

        # Combine the test and photo queries and normalize again
        search_features = text_features + query_photo_features * photo_weight
        search_features /= search_features.norm(dim=-1, keepdim=True)

        # Find the best match
        best_photo_ids = find_best_matches(search_features, photo_features, photo_ids, 10)
    else:
        # Display the results
        print("Test search result")
        best_photo_ids = search_unslash(query_text, photo_features, photo_ids, 10)

    return best_photo_ids


with gr.Blocks() as app:
    with gr.Row():
        gr.Markdown(
            """
            # CLIP Image Search Engine!
            ### Enter search query or/and input image to find the similar images from the database -
            """)

    with gr.Row(visible=True):
        with gr.Column():
            with gr.Row():
                search_text = gr.Textbox(value='', placeholder='Search..', label='Enter Your Query')

            with gr.Row():
                submit_btn = gr.Button("Submit", variant='primary')
                clear_btn = gr.ClearButton()

        with gr.Column():
            search_image = gr.Image(label='Upload Image or Select from results')

    with gr.Row(visible=True):
        output_images = gr.Gallery(allow_preview=False, label='Results.. ', info='',
                                   value=[], columns=5, rows=2)

        output_image_ids = gr.State([])


    def clear_data():
        return {
            search_image: None,
            output_images: None,
            search_text: None
        }


    clear_btn.click(clear_data, None, [search_image, output_images, search_text])


    def on_select(evt: gr.SelectData, output_image_ids):
        return {
            search_image: f"https://unsplash.com/photos/{output_image_ids[evt.index]}/download?w=100"
        }


    output_images.select(on_select, output_image_ids, search_image)


    def func_search(query, img):
        best_photo_ids = search_by_text_and_photo(query, img)
        img_urls = []
        for p_id in best_photo_ids:
            url = f"https://unsplash.com/photos/{p_id}/download?w=100"
            img_urls.append(url)

        valid_images = filter_invalid_urls(img_urls, best_photo_ids)

        return {
            output_image_ids: valid_images['image_ids'],
            output_images: valid_images['image_urls']
        }


    submit_btn.click(
        func_search,
        [search_text, search_image],
        [output_images, output_image_ids]
    )

'''
Launch the app
'''
app.launch()