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

def min_dim_to_size(img, size = 512):
    h, w = img.size
    ratio = size / max(h, w)
    h, w = map(lambda x: int(x * ratio), [h, w])
    return ( ratio ,img.resize((h, w)) )

#ratio_size = 512
#ratio, img_rs = min_dim_to_size(img, ratio_size)

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())
    if im is None:
        return list(map(lambda x: x["name"], images))

    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, ratio_size = 256 + 128):
    lex = Lexica(query=prompt).images()
    lex = lex[:limit_size]
    lex = list(map(lambda ele: Image.open(
                    requests.get(ele, stream = True).raw
                    ), lex))
    lex = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], lex))
    return lex

with gr.Blocks("custom.css") as demo:
    with gr.Row():
        with gr.Column():
            with gr.Row():
                #inputs = gr.Textbox(label = 'Enter prompt to search Lexica.art')
                inputs = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=20, min_width = 256,
                placeholder="Enter prompt to search Lexica.art", elem_id="prompt")
                #gr.Slider(label='Number of images ', minimum = 4, maximum = 20, step = 1, value = 4)]
                text_button = gr.Button("Retrieve Images", elem_id="run_button")
            i = gr.Image(elem_id="result-image")

        with gr.Column():
            title = gr.Markdown(
                    value="### Click on a Image in the gallery to select it, and the grid order will change",
                    visible=True,
                    elem_id="selected_model",
                )
            outputs = gr.Gallery(lable='Output gallery', elem_id="gallery",).style(grid=3,height=768,
                    allow_preview=False, label = "retrieve Images")
            #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)