#Importing libraries
import chromadb
import time
from transformers import CLIPModel, CLIPProcessor
import gradio as gr
from sklearn.metrics.pairwise import cosine_similarity
import torch
import numpy
from PIL import Image

#create the client
client = chromadb.Client()
collection = client.create_collection('image_collection')

#define the model and the proccessor
model = CLIPModel.from_pretrained('openai/clip-vit-base-patch32')
processor = CLIPProcessor.from_pretrained  ('openai/clip-vit-base-patch32')

image_paths = [
    'image-01.jpg',
    'image-02.jpg',
    'image-03.jpeg'
]

images = [Image.open(image_path) for image_path in image_paths]
inputs = processor(images = images, return_tensors='pt', padding=True)

#Generate the embeddings
start_time = time.time()
with torch.no_grad():
    image_embeddings = model.get_image_features(**inputs).numpy()

image_embeddings = [embedding.tolist() for embedding in image_embeddings]

end_time = time.time()
ingestion_time = end_time - start_time

collection.add(
    embeddings = image_embeddings,
    metadatas = [{'image':image} for image in image_paths],
    ids = [str(i) for i in range(len(image_paths))]
)

print(f'image ingestion time {ingestion_time:.4f} seconds')
def calculate_accuracy(image_embedding, query_embedding):
    similarity = cosine_similarity([image_embedding], [query_embedding])[0][0]
    return similarity

def search_image(query):
    if not query.strip():
        return None, 'Ooopsy, you forgot to input something?, please try again ☠👻'

    print(f'\nQuery: {query}')
    start_query = time.time()
    inputs = processor(text = query, return_tensors='pt', padding=True)
    with torch.no_grad():
        query_embeding = model.get_text_features(**inputs).numpy()

    query_embeding = query_embeding.tolist()
    end_query = time.time()
    query_time = end_query - start_query

    result = collection.query(query_embeddings = query_embeding, n_results=1)


    result_image_path = result['metadatas'][0][0]['image']
    result_image_index = int(result['ids'][0][0])
    matched_image_embedding = image_embeddings[result_image_index]

    accuracy_score = calculate_accuracy(matched_image_embedding, query_embeding[0])

    result_image = Image.open(result_image_path)
    file_name = result_image_path.split('/')[-1]

    return result_image, f'Accuracy score: {accuracy_score:.4f}\nQuery Time: {query_time:.4f} seconds', file_name

queries = [
    'A group of polar bears',
    'A famous landmark in paris',
    'A hot pizza fresh from the oven',
    'food',
    'A place',
    'A structure in Europe',
    'Animals'
]

def populate_queries(suggested_querries):
    return suggested_querries

#defining a gradio interface
with gr.Blocks() as gr_interface:
    gr.Markdown('# This is an image retrieval application Made by Allan Munene🤗😁🤑')

    with gr.Row():
        with gr.Column():
            gr.Markdown(f'***Image Ingestion Time***: {ingestion_time:.4f} seconds')
            gr.Markdown('## Input Panel')
            custom_query = gr.Textbox(placeholder = 'Input your query here', label='What are you looking for?')
            with gr.Row():
                submit_button = gr.Button('Submit Query')
                cancel_button = gr.Button('Cancel')
            
            with gr.Row(elem_id='button-container'):
                for query in queries:
                    gr.Button(query).click(fn=lambda q=query: q, outputs=custom_query)

    
        with gr.Column():
            gr.Markdown('Output Image')
            Output_image = gr.Image(type='pil', label='Result_image')

            accuracy = gr.Textbox(label='Performance')
            f_name = gr.Textbox(label='File Name')

        submit_button.click(fn=search_image, inputs=custom_query, outputs=[Output_image, accuracy, f_name])
        cancel_button.click(fn=lambda: (None, ''), outputs=[Output_image, accuracy, f_name])
        
gr_interface.launch(share=True)