import gradio as gr
import pandas as pd
import requests
import os
import numpy as np
import re
from tqdm import tqdm
from time import sleep
from PIL import Image
import requests
from io import BytesIO
from datasets import Dataset, load_dataset
import json
import cv2
import pathlib

import imagehash

MAX_MODEL_NUM = 300

'''
Yntec digplay
'''

hf_civital_image_info_dataset = load_dataset("svjack/hf_civital_image_info")
hf_civital_image_info_df = hf_civital_image_info_dataset["train"].to_pandas()

def gen_interface(model_name, max_times = 3):
    times = 0
    gr_model_interface = None
    while gr_model_interface is None and times < max_times:
        try:
            gr_model_interface = gr.load("models/{}".format(model_name),live=True)
        except:
            print("error {} times {}".format(model_name, times))
            sleep(2)
        times += 1
    return gr_model_interface

#gr_model_interface.title
def toImgOpenCV(imgPIL): # Conver imgPIL to imgOpenCV
    i = np.array(imgPIL) # After mapping from PIL to numpy : [R,G,B,A]
                         # numpy Image Channel system: [B,G,R,A]
    red = i[:,:,0].copy(); i[:,:,0] = i[:,:,2].copy(); i[:,:,2] = red;
    return i;

def toImgPIL(imgOpenCV): return Image.fromarray(cv2.cvtColor(imgOpenCV, cv2.COLOR_BGR2RGB));

def jpg_val_to_img(jpg_bytes):
    img_buf = np.frombuffer(jpg_bytes, np.uint8)
    img = cv2.imdecode(img_buf, cv2.IMREAD_UNCHANGED)
    return toImgPIL(img)


model_list = hf_civital_image_info_df["hf_repo_id"].drop_duplicates().values.tolist()
model_interface_list = []
for model_name in tqdm(model_list):
    gr_model_interface = gen_interface(model_name)
    if gr_model_interface is not None:
        model_interface_list.append(gr_model_interface)
        if len(model_interface_list) >= MAX_MODEL_NUM:
            break

def get_civital_iframe(url, width = 1400, height = 768, as_html = True, visible = False):
    html= '''
            <div style="justify-content: center; display: flex;">
                    <iframe
                    src="{}"
                    frameborder="0"
                    width="{}"
                    height="{}"
                    ></iframe>
            </div>
        '''.format(url, width, height)
    if as_html:
        html = gr.HTML(html, visible = visible)
    return html

def get_info_by_interface(gr_interface, model_interface_list = model_interface_list):
    #### out: (gr_interface, civital_url, civital_info)
    if hasattr(gr_interface, "app"):
        civital_url = hf_civital_image_info_df[
            hf_civital_image_info_df["hf_repo_id"] == gr_interface.title
        ]["civital_url"].iloc[0]
        civital_info = hf_civital_image_info_df[
            hf_civital_image_info_df["hf_repo_id"] == gr_interface.title
        ][["prompt", "image"]].values.tolist()
        return gr_interface ,civital_url, civital_info
    else:
        civital_url = hf_civital_image_info_df[
            hf_civital_image_info_df["hf_repo_id"] == gr_interface
        ]["civital_url"].iloc[0]
        civital_info = hf_civital_image_info_df[
            hf_civital_image_info_df["hf_repo_id"] == gr_interface
        ][["prompt", "image"]].values.tolist()
        return list(filter(lambda x:x.title == gr_interface, model_interface_list))[0] ,civital_url, civital_info

def read_image_from_url(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    return img

def image_click(images, evt: gr.SelectData, gr_interface_value,
    ):
    img_selected = images[evt.index]
    #print(img_selected)
    im_data = img_selected["name"]
    im = Image.open(im_data)
    im_hash = imagehash.average_hash(
        im, hash_size = 1024
    )
    min_diff = int(1e10)
    #print(-1)
    repo_card_im_dict = dict(
    get_info_by_interface(gr_interface_value)[2]
    )
    min_repo_name = ""
    for idx ,(repo_name, repo_card_image) in enumerate(repo_card_im_dict.items()):
        repo_img = jpg_val_to_img(repo_card_image["bytes"])
        repo_img_hash = imagehash.average_hash(
            repo_img, hash_size = 1024
        )
        diff = im_hash - repo_img_hash
        if diff < min_diff:
            min_diff = diff
            min_repo_name = repo_name
        #print(idx)
    prompt = min_repo_name
    return prompt
    #return prompt, im

def try_repo_act_func(civital_url, show_civital_button):
    repo_html_iframe_hide = get_civital_iframe(civital_url, visible = True if show_civital_button == "Show Civital Page" else False)
    return repo_html_iframe_hide,  gr.Button("Hide Civital Page" if show_civital_button == "Show Civital Page" else "Show Civital Page")

with gr.Blocks(
    css = '''
    .header img {
          float: middle;
          width: 33px;
          height: 33px;
        }
        .header h1 {
          top: 18px;
          left: 10px;
        }
    '''
) as demo:
    gr.HTML(
        '''
        <center>
        <div class="header">
        <h1 class = "logo"> <img src="https://huggingface.co/spaces/svjack/Civital-Stable-Diffusion-HF/resolve/main/logo.png" alt="logo" />  🤗 Civital Model on Huggingface </h1>
        </center>
        '''
    )

    with gr.Row():
        with gr.Column():
            with gr.Row():
                hf_model_dropdown = gr.Dropdown(label = "🤗 Hf model",
                    choices=list(map(lambda x: x.title, model_interface_list)),
                    value=list(map(lambda x: x.title, model_interface_list))[0],)
            with gr.Column():
                with gr.Row():
                    civital_prompt = gr.Textbox(label = "🤭 Civital Prompt (Click from 👉 right gallery to get them, and You can edit ✍️ yourself)",
                        interactive = True,
                    )
                    gen_button = gr.Button(label = "Generate")
                hf_image = gr.Image(label = "🤭 Image generate by 🤗 Huggingface", height = 768)

        with gr.Column():
            civital_info_gallery = gr.Gallery(
                pd.Series(
                get_info_by_interface(hf_model_dropdown.value)[2]
                ).sample(n = min(len(get_info_by_interface(hf_model_dropdown.value)[2]), 30)).map(lambda t2: t2[1]).map(lambda x: x["bytes"]).map(jpg_val_to_img).values.tolist(),
                height = 1024,
                label = "🖱️👇 ➡️ 👈 Civital image samples",
                object_fit = "contain"
            )

    with gr.Row():
        with gr.Column():
            try_repo_button = gr.Button("Show Civital Page")
            civital_iframe_html = get_civital_iframe(
                get_info_by_interface(hf_model_dropdown.value)[1]
            )

    hf_model_dropdown.change(
        lambda x: pd.Series(
        get_info_by_interface(x)[2]
        ).sample(n = min(len(get_info_by_interface(x)[2]), 30)).map(lambda t2: t2[1]).map(lambda x: x["bytes"]).map(jpg_val_to_img).values.tolist(),
        hf_model_dropdown,
        civital_info_gallery
    )
    hf_model_dropdown.change(
        lambda _: (gr.Button("Show Civital Page"), gr.HTML(visible = False)),
        None,
        [try_repo_button, civital_iframe_html]
    )

    civital_info_gallery.select(
            image_click,
            [civital_info_gallery, hf_model_dropdown],
            civital_prompt
    )
    gen_button.click(lambda hf_model_name, text_prompt:
            get_info_by_interface(hf_model_name)[0](text_prompt),
            [hf_model_dropdown, civital_prompt],
            hf_image
    )

    try_repo_button.click(
        lambda hf_model_name, button: try_repo_act_func(
            get_info_by_interface(hf_model_name)[1]
            , button),
        [hf_model_dropdown, try_repo_button],
        [civital_iframe_html, try_repo_button]
    )

demo.launch(show_api = False)