from vis_common import *
import vis_utils as v_uts
import io_utils as io_uts

import os
from pickle import FALSE
import gradio as gr
from functools import partial
import json

# install gradio of 3.14
os.system("echo $BYTED_HOST_IP")

# Load the dataset change to your local path
root = "/home/mudehui/ChatEdit"
#prompt_version = "prompt_0_sd"
#prompt_version = "prompt_0_hd"
#prompt_version = "prompt_1_sd"
#prompt_version = "prompt_1_hd"
prompt_version = "prompt_0_rewrited_sd"

def load_json(file, existing_data=[]):
    if not os.path.exists(file):
        empty = {}
        return empty
    with open(file, "r") as f:
        stats = json.load(f)

    results = {name: score for name, score in stats.items() \
        if name not in existing_data}
    return results

all_items = f"{root}/full_val.jsonl"
all_samples = io_uts.load_jsonl(all_items)
all_samples = {f"{i:03}":all_samples[i] for i in range(len(all_samples))}

votes = {}
def update(name, picture_name, vote, start_idx=0, end_idx=1000):
    record_file = f"./output/{prompt_version}/{name}.json"
    v_uts.mkdir("", record_file)
    start_idx, end_idx = int(start_idx), int(end_idx)
    end_idx = min(end_idx, len(all_samples) - 1)
    items = list(all_samples.items())[start_idx:end_idx]
    label_samples = {name:prompt for name, prompt in items}

    if name == "":
        new_picture = None
        picture_name = None
        description = None
        message = "Please enter your lark username"

    elif picture_name in label_samples.keys() and vote is None:
        new_picture = None
        picture_name = None
        description = None
        message = "Please make selections! Click Next to continue..."

    else:
        # Read record
        existing_data = load_json(record_file)

        # Save record
        if (picture_name in label_samples.keys()):
            sample = label_samples[picture_name]
            sample["vote"] = vote
            existing_data[picture_name] = sample
            with open(record_file, "w") as f:
                json.dump(existing_data, f, indent=2)

        # Find Next example
        all_remaining = {}
        for i, name in enumerate(label_samples.keys()):
            if name in existing_data:
                continue
            else:
                all_remaining[name] = label_samples[name]

        if len(all_remaining) > 0:
            new_sample = list(all_remaining.items())[0]
            picture_name, data = new_sample
            description = f"input: {data['input']}<br>output: {data['output']}<br>edit: {data['edit']}"
            new_picture = f"{root}/{prompt_version}/{picture_name}.png"
            message = f"{len(all_remaining)} exmaples remaining"
        else:
            new_picture = None
            picture_name = None
            description = None
            message = "You have finished all exmaples! Thank you!"

    outputs = [new_picture, picture_name, message, description]
    print(outputs)
    return tuple(outputs)


with gr.Blocks() as demo:

    gr.Markdown("""
         - 输入用户名, 开始结束index，点击Next按钮开始, 你正在评价 {prompt}"
    """.format(prompt=prompt_version))
    with gr.Row():
        with gr.Column():
            picture_name = gr.Textbox(visible=FALSE)
            picture = gr.Image(label=f"Input Image from ")

        with gr.Column():
            name = gr.Textbox(label="User Name (enter and click Next to start)")
            start_idx = gr.Textbox(label="Start Index (max 292)", default="0")
            end_idx = gr.Textbox(label="End Index (max 292)", default="1000")
            message = gr.Markdown()
            description = gr.Markdown()
            vote = gr.Radio([
                ('1: Totally not related ', 1), 
                ('2: Not follow edit, there is some/little relation between the two images.', 2), 
                ('3: OK Pair data, not follow edit, image pair need some edit effort [flip etc.] to construct a good edit pair.', 3),
                ('4: Good pair data, can modify the instruction to form a good triplet', 4), 
                ('5: Perfectly follows the edit instruction.', 5)
            ], label="Score", min_width=400)
            greet_btn = gr.Button("Next")
    greet_btn.click(fn=update, 
                    inputs=[name,picture_name,vote, start_idx, end_idx], 
                    outputs=[picture,picture_name,message,description])

demo.queue(max_size=4)
demo.launch(share=True)