import gradio as gr
from Lex import *
'''
lex = Lexica(query="man woman fire snow").images()
'''
from PIL import Image
import imagehash
import requests

def image_click(images, evt: gr.SelectData):
    img_selected = images[evt.index]
    return images[evt.index]['name']

def swap_gallery(im, images):
    #### name data is_file
    #print(images[0].keys())

    im_hash = imagehash.average_hash(Image.fromarray(im))
    t2_list = sorted(images, key = lambda imm:
    imagehash.average_hash(Image.open(imm["name"])) - im_hash, reverse = False)
    return list(map(lambda x: x["name"], t2_list))


def lexica(prompt, limit_size = 128):
    lex = Lexica(query=prompt).images()
    lex = lex[:limit_size]
    lex = list(map(lambda ele: Image.open(
                    requests.get(ele, stream = True).raw
                    ), lex))
    return lex

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Tabs():
            inputs = gr.Textbox(label = 'Enter prompt to search Lexica.art')
            text_button = gr.Button("Retrieve Images")
            #gr.Slider(label='Number of images ', minimum = 4, maximum = 20, step = 1, value = 4)]
        with gr.Tabs():
            i = gr.Image()
            outputs = gr.Gallery(lable='Output gallery').style(grid=3,height=768,
                allow_preview=False)
            #gr.Dataframe(label='prompts for corresponding images')]
            
        #outputs.select(image_click, outputs, i, _js="(x) => x.splice(0,x.length)")    
        outputs.select(image_click, outputs, i,)
        
        i.change(
                fn=swap_gallery,
                inputs=[i, outputs],
                outputs=outputs,
                queue=False
                )
        #### gr.Textbox().submit().success()

        text_button.click(lexica, inputs=inputs, outputs=outputs)

demo.launch("0.0.0.0" ,debug = True)