from .utils import *
from .vote_utils import (
    upvote_last_response_ie as upvote_last_response,
    downvote_last_response_ie as downvote_last_response,
    flag_last_response_ie as flag_last_response,
    leftvote_last_response_iem as leftvote_last_response,
    rightvote_last_response_iem as rightvote_last_response,
    tievote_last_response_iem as tievote_last_response,
    bothbad_vote_last_response_iem as bothbad_vote_last_response,
    share_click_iem as share_click,
    generate_ie,
    generate_ie_museum,
    generate_iem,
    generate_iem_museum,
    generate_iem_annoy,
    generate_iem_annoy_museum,
    share_js
)
from functools import partial

def build_side_by_side_ui_anony_ie(models):
    notice_markdown = """
# ⚔️  GenAI-Arena ⚔️ : Benchmarking Visual Generative Models in the Wild
| [Code](https://huggingface.co/spaces/TIGER-Lab/GenAI-Arena/tree/main) | [Paper](https://arxiv.org/abs/2406.04485v1) | [Dataset](https://huggingface.co/datasets/TIGER-Lab/GenAI-Bench) |

## 📜 Rules
- Edit with two selected models side-by-side and vote!
- Upload a source image that you want to edit.
- Input source prompt, target prompt and instruct prompt.
- Wait to see the results after edition.
- Click "New Round" to start a new round.
- Vote won't be counted if model identity is revealed during generation.
- The model coulmodel_selector_left.changed output a totally black image or noise. It's not a bug but the failure case of the model.

## ⚠️ Data Collection Consent
- Your input and votes will be collected for research purposes only. 
- By using this service, you agree to the collection of your input and votes for research purposes. 
- Your data will be anonymized and will not be used for commercial purposes.

## 🏆 Arena Elo 
Find out who is the 🥇conditional image edition models!

## 👇 Editing now!

"""

    model_list = models.model_ie_list

    state0, state1 = gr.State(), gr.State()
    gen_func = partial(generate_iem_annoy, models.generate_image_ie_parallel_anony)
    gen_func_random = partial(generate_iem_annoy_museum, models.generate_image_ie_museum_parallel_anony)
    gr.Markdown(notice_markdown, elem_id="notice_markdown")

    with gr.Group(elem_id="share-region-anony"):
        with gr.Accordion("🔍 Expand to see all Arena players", open=False):
            model_description_md = get_model_description_md(model_list)
            gr.Markdown(model_description_md, elem_id="model_description_markdown")
        with gr.Row():
            with gr.Column():
                    chatbot_left = gr.Image(width=512, type="pil")
            with gr.Column():
                    chatbot_right = gr.Image(width=512, type="pil")

        with gr.Row():
            with gr.Column():
                model_selector_left =gr.Markdown("", visible=False)
            with gr.Column():
                model_selector_right = gr.Markdown("", visible=False)
        with gr.Row():
            slow_warning = gr.Markdown("", elem_id="notice_markdown")

        with gr.Row():
            leftvote_btn = gr.Button(
                value="👈  A is better", visible=False, interactive=False
            )
            rightvote_btn = gr.Button(
                value="👉  B is better", visible=False, interactive=False
            )
            tie_btn = gr.Button(value="🤝  Tie", visible=False, interactive=False)
            bothbad_btn = gr.Button(
                value="👎  Both are bad", visible=False, interactive=False
            )

    with gr.Row():

        textbox_source = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your source prompt",
            elem_id="input_box_s",
        )
        textbox_target = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your target prompt",
            elem_id="input_box_t",
        )
        textbox_instruct = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your instruct prompt",
            elem_id="input_box_t",
        )

    with gr.Row():
        source_image = gr.Image(type="pil")
        send_btn = gr.Button(value="Send", variant="primary", scale=0)
        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)

    with gr.Row() as button_row:
        clear_btn = gr.Button(value="🎲 New Round", interactive=False)
        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
        # share_btn = gr.Button(value="📷  Share")

    #gr.Markdown(acknowledgment_md, elem_id="ack_markdown")

    gr.Examples(
        examples=[
            ["a bowl of strawberries", "a bowl of oranges", "change strawberries to oranges",
             os.path.join("./examples", "strawberries.jpg")],
            ["a pig is eating an ice cream", "a rabbit is eating an ice cream", "change pig to rabbit",
             os.path.join("./examples", "pig.jpg")],
            ["a rubber duck in a swimming pool", "a rubber duck with a hat in a swimming pool", "add a hat to the duck",
             os.path.join("./examples", "duck.jpg")],
            ["a photo of a cat", "a photo of a mouse", "change cat to mouse",
             os.path.join("./examples", "cat.jpeg")]],
        inputs=[textbox_source, textbox_target, textbox_instruct, source_image])

    # random_models = random.sample([model for model in model_list], 2)
    # model_selector_left =gr.Markdown(random_models[0])
    # model_selector_right =gr.Markdown(random_models[1])
    
    btn_list = [leftvote_btn, rightvote_btn, tie_btn, bothbad_btn, regenerate_btn, clear_btn]

    send_btn.click(
        gen_func,
        inputs=[
            state0,
            state1,
            textbox_source, 
            textbox_target, 
            textbox_instruct, 
            source_image, 
            model_selector_left, 
            model_selector_right
        ],
        outputs=[state0, state1, chatbot_left, chatbot_right, model_selector_left, model_selector_right],
        api_name="send_btn_side_by_side_anony"
    ).then(
        enable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    draw_btn.click(
        gen_func_random,
        inputs=[
            state0,
            state1,
            model_selector_left, 
            model_selector_right
        ],
        outputs=[state0, state1, chatbot_left, chatbot_right, 
                source_image, textbox_source, textbox_target, textbox_instruct, 
                model_selector_left, model_selector_right],
        api_name="draw_btn_side_by_side_anony"
    ).then(
        enable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    regenerate_btn.click(
        gen_func,
        inputs=[
            state0,
            state1,
            textbox_source, 
            textbox_target, 
            textbox_instruct, 
            source_image, 
            model_selector_left, 
            model_selector_right
        ],
        outputs=[state0, state1, chatbot_left, chatbot_right, model_selector_left, model_selector_right],
        api_name="regenerate_btn_side_by_side_anony"
    ).then(
        enable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    
    clear_btn.click(
        clear_history_side_by_side_ie_anony, 
        inputs=None, 
        outputs=[state0, state1, textbox_source, textbox_target, textbox_instruct, 
                 source_image, chatbot_left, chatbot_right,
                 model_selector_left, model_selector_right], 
        api_name="clear_btn_side_by_side_anony"
    ).then(
        disable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    vote_btns = [leftvote_btn, rightvote_btn, tie_btn, bothbad_btn]
    input_list = [textbox_source, textbox_target, source_image, textbox_instruct]
    dummy_left_model = gr.State("")
    dummy_right_model = gr.State("")
    leftvote_btn.click(
        leftvote_last_response, 
        inputs=[state0, state1, dummy_left_model, dummy_right_model],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )
    rightvote_btn.click(
        rightvote_last_response,
        inputs=[state0, state1, dummy_left_model, dummy_right_model],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )
    tie_btn.click(
        tievote_last_response,
        inputs=[state0, state1, dummy_left_model, dummy_right_model],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )
    bothbad_btn.click(
        bothbad_vote_last_response,
        inputs=[state0, state1, dummy_left_model, dummy_right_model],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )



    

def build_side_by_side_ui_named_ie(models):
    notice_markdown = """
# ⚔️ GenAI-Arena ⚔️ : Benchmarking Visual Generative Models in the Wild
| [Code](https://huggingface.co/spaces/TIGER-Lab/GenAI-Arena/tree/main) | [Paper](https://arxiv.org/abs/2406.04485v1) | [Dataset](https://huggingface.co/datasets/TIGER-Lab/GenAI-Bench) | 

## 📜 Rules
- Edit with two selected models side-by-side and vote!
- Upload a source image that you want to edit.
- Input source prompt, target prompt and instruct prompt.
- Wait to see the results after edition.
- Click "Clear history" to start a new round.
- The model could output a totally black image or noise. It's not a bug but the failure case of the model.

## ⚠️ Data Collection Consent
- Your input and votes will be collected for research purposes only. 
- By using this service, you agree to the collection of your input and votes for research purposes. 
- Your data will be anonymized and will not be used for commercial purposes.

## 🤖 Choose two models to compare
"""

    model_list = models.model_ie_list

    state0, state1 = gr.State(), gr.State()
    gen_func = partial(generate_iem, models.generate_image_ie_parallel)
    gen_func_random = partial(generate_iem_museum, models.generate_image_ie_museum_parallel)
    gr.Markdown(notice_markdown, elem_id="notice_markdown")

    with gr.Group(elem_id="share-region-named"):
        with gr.Row():
            with gr.Column():
                model_selector_left = gr.Dropdown(
                    choices=model_list,
                    value=model_list[0] if len(model_list) > 0 else "",
                    interactive=True,
                    show_label=False,
                    container=False,
                    )
            with gr.Column():
                model_selector_right = gr.Dropdown(
                    choices=model_list,
                    value=model_list[1] if len(model_list) > 1 else "",
                    interactive=True,
                    show_label=False,
                    container=False,
                    )
        with gr.Row():
            with gr.Accordion("🔍 Expand to see all model descriptions", open=False):
                model_description_md = get_model_description_md(model_list)
                gr.Markdown(model_description_md, elem_id="model_description_markdown")

        with gr.Row():
            with gr.Column():
                chatbot_left = gr.Image(width=512, label = "Model A")
            with gr.Column():
                chatbot_right = gr.Image(width=512, label = "Model B")

        with gr.Row():
            leftvote_btn = gr.Button(
                value="👈  A is better", visible=False, interactive=False
            )
            rightvote_btn = gr.Button(
                value="👉  B is better", visible=False, interactive=False
            )
            tie_btn = gr.Button(value="🤝  Tie", visible=False, interactive=False)
            bothbad_btn = gr.Button(
                value="👎  Both are bad", visible=False, interactive=False
            )

    with gr.Row():

        textbox_source = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your source prompt",
            elem_id="input_box_s",
        )
        textbox_target = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your target prompt",
            elem_id="input_box_t",
        )
        textbox_instruct = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your instruct prompt",
            elem_id="input_box_t",
        )

    with gr.Row():
        source_image = gr.Image(type="pil")
        send_btn = gr.Button(value="Send", variant="primary", scale=0)
        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)

    with gr.Row() as button_row:
        clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
        # share_btn = gr.Button(value="📷  Share")

    #gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
    gr.Examples(
        examples=[
            ["a bowl of strawberries", "a bowl of oranges", "change strawberries to oranges",
             os.path.join("./examples", "strawberries.jpg")],
            ["a pig is eating an ice cream", "a rabbit is eating an ice cream", "change pig to rabbit",
             os.path.join("./examples", "pig.jpg")],
            ["a rubber duck in a swimming pool", "a rubber duck with a hat in a swimming pool", "add a hat to the duck",
             os.path.join("./examples", "duck.jpg")],
            ["a photo of a cat", "a photo of a mouse", "change cat to mouse",
             os.path.join("./examples", "cat.jpeg")]],
        inputs=[textbox_source, textbox_target, textbox_instruct, source_image])
    

    # model_selector_left.change(
    #     clear_history_side_by_side_ie,
    #     inputs=None,
    #     outputs=[state0, state1, textbox_source, textbox_target, textbox_instruct,
    #              source_image, chatbot_left, chatbot_right
    #              ],
    #     api_name="model_selector_left_side_by_side")
    # model_selector_right.change(clear_history_side_by_side_ie, inputs=None, outputs=[state0, state1, textbox_source, textbox_target, textbox_instruct, source_image, chatbot_left, chatbot_right], api_name="model_selector_right_side_by_side")


    btn_list = [leftvote_btn, rightvote_btn, tie_btn, bothbad_btn, regenerate_btn, clear_btn]

    send_btn.click(
        gen_func,
        inputs=[
            state0,
            state1,
            textbox_source, 
            textbox_target, 
            textbox_instruct, 
            source_image, 
            model_selector_left, 
            model_selector_right
        ],
        outputs=[state0, state1, chatbot_left, chatbot_right],
        api_name="send_btn_side_by_side"
    ).then(
        enable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    draw_btn.click(
        gen_func_random,
        inputs=[
            state0,
            state1,
            model_selector_left, 
            model_selector_right
        ],
        outputs=[state0, state1, chatbot_left, chatbot_right, 
                 source_image, textbox_source, textbox_target, textbox_instruct 
                ],
        api_name="draw_btn_side_by_side"
    ).then(
        enable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    regenerate_btn.click(
        gen_func,
        inputs=[
            state0,
            state1,
            textbox_source, 
            textbox_target, 
            textbox_instruct, 
            source_image, 
            model_selector_left, 
            model_selector_right
        ],
        outputs=[state0, state1, chatbot_left, chatbot_right],
        api_name="regenerate_btn_side_by_side"
    ).then(
        enable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    
    clear_btn.click(
        clear_history_side_by_side_ie, 
        inputs=None, 
        outputs=[state0, state1, textbox_source, textbox_target, textbox_instruct, 
                 source_image, chatbot_left, chatbot_right], 
        api_name="clear_btn_side_by_side"
        ).then(
        disable_buttons_side_by_side, 
        inputs=None,  
        outputs=btn_list 
    )
    vote_btns = [leftvote_btn, rightvote_btn, tie_btn, bothbad_btn]
    input_list = [textbox_source, textbox_target, source_image, textbox_instruct]
    leftvote_btn.click(
        leftvote_last_response, 
        inputs=[state0, state1, model_selector_left, model_selector_right],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )
    rightvote_btn.click(
        rightvote_last_response,
        inputs=[state0, state1, model_selector_left, model_selector_right],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )
    tie_btn.click(
        tievote_last_response,
        inputs=[state0, state1, model_selector_left, model_selector_right],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )
    bothbad_btn.click(
        bothbad_vote_last_response,
        inputs=[state0, state1, model_selector_left, model_selector_right],
        outputs=[model_selector_left, model_selector_right] + input_list + vote_btns,
    )



def build_single_model_ui_ie(models, add_promotion_links=False):
    promotion = (
        """
- | [Code](https://huggingface.co/spaces/TIGER-Lab/GenAI-Arena/tree/main) | [Paper](https://arxiv.org/abs/2406.04485v1) | [Dataset](https://huggingface.co/datasets/TIGER-Lab/GenAI-Bench) |
"""
        if add_promotion_links
        else ""
    )

    notice_markdown = f"""
# 🏔️ Play with Image Edition Models
{promotion}

## 🤖 Choose any model to edit image
"""
    gr.Markdown(notice_markdown, elem_id="notice_markdown")

    state = gr.State()
    gen_func = partial(generate_ie, models.generate_image_ie)
    gen_func_random = partial(generate_ie_museum, models.generate_image_ie_museum)
    model_list = models.model_ie_list

    with gr.Row(elem_id="model_selector_row"):
        model_selector = gr.Dropdown(
            choices=model_list,
            value=model_list[0] if len(model_list) > 0 else "",
            interactive=True,
            show_label=False
        )
    with gr.Row():
        with gr.Accordion(
            "🔍 Expand to see all model descriptions",
            open=False,
            elem_id="model_description_accordion",
        ):
            model_description_md = get_model_description_md(model_list)
            gr.Markdown(model_description_md, elem_id="model_description_markdown")

    with gr.Row():

        textbox_source = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your source prompt",
            elem_id="input_box_s",
        )
        textbox_target = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your target prompt",
            elem_id="input_box_t",
        )
        textbox_instruct = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your instruct prompt",
            elem_id="input_box_t",
        )

    with gr.Row():
        source_image = gr.Image(type="pil")
        send_btn = gr.Button(value="Send", variant="primary", scale=0)
        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)

    with gr.Row():
        chatbot = gr.Image()

    with gr.Row() as button_row:
        upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
        downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
        flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
        clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)

    if add_promotion_links:
        #gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
        pass

    gr.Examples(
        examples=[
            ["a bowl of strawberries", "a bowl of oranges", "change strawberries to oranges",
             os.path.join("./examples", "strawberries.jpg")],
            ["a pig is eating an ice cream", "a rabbit is eating an ice cream", "change pig to rabbit",
             os.path.join("./examples", "pig.jpg")],
            ["a rubber duck in a swimming pool", "a rubber duck with a hat in a swimming pool", "add a hat to the duck",
             os.path.join("./examples", "duck.jpg")],
            ["a photo of a cat", "a photo of a mouse", "change cat to mouse",
             os.path.join("./examples", "cat.jpeg")]],
        inputs=[textbox_source, textbox_target, textbox_instruct, source_image])

    # Register listeners
    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
    
    model_selector.change(clear_history_ie, inputs=None, outputs=[state, textbox_source, textbox_target, textbox_instruct, source_image, chatbot], api_name="model_selector_single")

    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]

    send_btn.click(
        gen_func,
        inputs=[state, textbox_source, textbox_target, textbox_instruct, source_image, model_selector], 
        outputs=[state, chatbot], 
        api_name="send_btn_single",
        show_progress = "full"
    ).then(
        enable_buttons, 
        inputs=None,  
        outputs=btn_list 
    )
    draw_btn.click(
        gen_func_random,
        inputs=[state, model_selector], 
        outputs=[state, chatbot, source_image, textbox_source, textbox_target, textbox_instruct], 
        api_name="send_btn_single",
        show_progress = "full"
    ).then(
        enable_buttons, 
        inputs=None,  
        outputs=btn_list 
    )
    

    vote_btns = [upvote_btn, downvote_btn, flag_btn]
    model_inputs = [textbox_source, textbox_instruct, source_image, textbox_target]
    upvote_btn.click(
        upvote_last_response,
        inputs=[state, model_selector],
        outputs=model_inputs + vote_btns,
    )
    downvote_btn.click(
        downvote_last_response, 
        inputs=[state, model_selector],
        outputs=model_inputs + vote_btns,
    )
    flag_btn.click(
        flag_last_response,
        inputs=[state, model_selector],
        outputs=model_inputs + vote_btns,
    )
    regenerate_btn.click(
        gen_func,
        inputs=[state, textbox_source, textbox_target, textbox_instruct, source_image, model_selector], 
        outputs=[state, chatbot], 
        api_name="regenerate_btn_single",
        show_progress = "full"
    ).then(
        enable_buttons, 
        inputs=None,  
        outputs=btn_list 
    )
    clear_btn.click(
        clear_history_ie, 
        inputs=None, 
        outputs=[state, textbox_source, textbox_target, textbox_instruct, source_image, chatbot], 
        api_name="clear_history_single",
        show_progress="full"
    ).then(
        disable_buttons, 
        inputs=None,  
        outputs=btn_list 
    )