from .utils import *
from .vote_utils import (
    leftvote_last_response_igm as leftvote_last_response,
    rightvote_last_response_igm as rightvote_last_response,
    bothbad_vote_last_response_igm as bothbad_vote_last_response,
    leftvote_last_response_igm_anony as leftvote_last_response_anony,
    rightvote_last_response_igm_anony as rightvote_last_response_anony,
    bothbad_vote_last_response_igm_anony as bothbad_vote_last_response_anony,
)
# from functools import partial
        
def build_side_by_side_ui_anony(models):
    notice_markdown = \
        """
        # ⚔️  IDEA-Bench ⚔️ : How Far are Generative Models from Professional Designing?
        | [GitHub](https://github.com/ali-vilab/IDEA-Bench) | [Paper](https://arxiv.org/abs/2412.11767) | [Dataset](https://huggingface.co/datasets/ali-vilab/IDEA-Bench) | 
        
        ## 📜 Rules
        - Choose the better one from two anonymous models.
        - Click "New Round" to start a new round.
        - After the voting ends, the model name will be displayed and the voting selection cannot be changed.

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

        ## 🏆 Arena Elo 
        Find out the best model for professional level image processing tasks! Welcome to upload your own model generation results!

        ## 👇 Voting now!

        """

    state0 = gr.State()
    state1 = gr.State()
    state_models = gr.State(models)  # 使用 State 存储模型
    
    gr.Markdown(notice_markdown, elem_id="notice_markdown")

    with gr.Group(elem_id="share-region-anony"):
        # Display model description
        with gr.Accordion("🔍 Expand to see all Arena players", open=False):
            model_description_md = models.get_model_info_md()
            gr.Markdown(model_description_md, elem_id="model_description_markdown")
        
        # Display the input prompt as non-editable text
        with gr.Accordion("✏️ The input prompt:", open=True):
            prompt = gr.Markdown( value="",  # Display the prompt here
                elem_id="prompt"
            )
        
        # Display input images
        with gr.Row():
            input_images = gr.Gallery(
                value=[],
                columns=5, 
                label="Input Images",
                show_label=True,
                elem_id="input_images",
                preview=True,
                visible=True
            )

        with gr.Row():
            with gr.Column():
                chatbot_left = gr.Gallery(
                    value=[],
                    columns=3,
                    label="Model A",
                    show_label=True,
                    preview=True
                )
            with gr.Column():
                chatbot_right = gr.Gallery(
                    value=[],
                    columns=3,
                    label="Model B",
                    show_label=True,
                    preview=True
                )

        with gr.Row():
            with gr.Column():
                model_demo_left = gr.Markdown("", visible=False)
            with gr.Column():
                model_demo_right = gr.Markdown("", visible=False)

        with gr.Row():
            slow_warning = gr.Markdown("", elem_id="notice_markdown")

        with gr.Row(elem_classes=["btn-row"]):
            with gr.Column():
                leftvote_btn = gr.Button(
                    value="Prefer 👈", 
                    visible=True, 
                    interactive=False,
                    elem_classes=["vote-btn"],
                )
            with gr.Column():
                rightvote_btn = gr.Button(
                    value="Prefer 👉", 
                    visible=True, 
                    interactive=False,
                    elem_classes=["vote-btn"],
                )
        
        with gr.Row(elem_classes=["btn-row"]):
            with gr.Column():
                bothbad_btn = gr.Button(
                    value="Both are bad 👎", 
                    visible=True, 
                    interactive=False,
                    elem_classes=["vote-btn"],
                )
            with gr.Column():
                refresh_btn = gr.Button(
                    value="New Round 🎲", 
                    interactive=True,
                    elem_classes=["vote-btn"],
                )

    btn_list = [leftvote_btn, rightvote_btn, bothbad_btn, refresh_btn]
    
    refresh_btn.click(
        refresh_side_by_side_anony, 
        inputs=[state_models], 
        outputs=[state0, state1, prompt, input_images, chatbot_left, chatbot_right, leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right], 
        api_name="refresh_btn_annony"
    )

    leftvote_btn.click(
        leftvote_last_response_anony,
        inputs=[state0, state1],
        outputs=[leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right]
    )

    rightvote_btn.click(
        rightvote_last_response_anony,
        inputs=[state0, state1],
        outputs=[leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right]
    )

    bothbad_btn.click(
        bothbad_vote_last_response_anony,
        inputs=[state0, state1],
        outputs=[leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right]
    )
    
    
def build_side_by_side_ui_named(models):
    notice_markdown = \
        """
        # ⚔️  IDEA-Bench ⚔️ : How Far are Generative Models from Professional Designing?
        | [GitHub](https://github.com/ali-vilab/IDEA-Bench) | [Paper](https://arxiv.org/abs/2412.11767) | [Dataset](https://huggingface.co/datasets/ali-vilab/IDEA-Bench) | 
        
        ## 📜 Rules
        - Choose the better one from two anonymous models.
        - Click "New Round" to start a new round.
        - After the voting ends, the model name will be displayed and the voting selection cannot be changed.

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

        ## 🏆 Arena Elo 
        Find out the best model for professional level image processing tasks! Welcome to upload your own model generation results!

        ## 👇 Voting now!

        """
    
    model_list = models.get_name_list()

    state0 = gr.State()
    state1 = gr.State()
    state_models = gr.State(models)  # 使用 State 存储模型

    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,
                    allow_custom_value=True
                    )
            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,
                    allow_custom_value=True
                    )
            # with gr.Column():
            #     gen_btn = gr.Button(
            #         value="Generate 🔄", visible=True, interactive=True
            #     )
        # Display model description
        with gr.Accordion("🔍 Expand to see all Arena players", open=False):
            model_description_md = models.get_model_info_md()
            gr.Markdown(model_description_md, elem_id="model_description_markdown")
        
        # Display the input prompt as non-editable text
        with gr.Accordion("✏️ The input prompt:", open=True):
            prompt = gr.Markdown(
                value="",  # Display the prompt here
                elem_id="prompt"
            )
        
        # Display input images
        with gr.Row():
            input_images = gr.Gallery(
                value=[],
                columns=5, 
                label="Input Images",
                show_label=True,
                elem_id="input_images",
                preview=True,
                visible=True
            )

        with gr.Row():
            with gr.Column():
                chatbot_left = gr.Gallery(
                    value=[],
                    columns=3,
                    label="Model A",
                    show_label=True,
                    preview=True
                )
            with gr.Column():
                chatbot_right = gr.Gallery(
                    value=[],
                    columns=3,
                    label="Model B",
                    show_label=True,
                    preview=True
                )

        with gr.Row():
            with gr.Column():
                model_demo_left = gr.Markdown("", visible=False)
            with gr.Column():
                model_demo_right = gr.Markdown("", visible=False)

        with gr.Row():
            slow_warning = gr.Markdown("", elem_id="notice_markdown")

        with gr.Row(elem_classes=["btn-row"]):
            with gr.Column():
                leftvote_btn = gr.Button(
                    value="Prefer 👈", 
                    visible=True, 
                    interactive=False, 
                    elem_classes=["vote-btn"],
                )
            with gr.Column():
                rightvote_btn = gr.Button(
                    value="Prefer 👉", 
                    visible=True, 
                    interactive=False, 
                    elem_classes=["vote-btn"],
                )
        
        with gr.Row(elem_classes=["btn-row"]):
            with gr.Column():
                bothbad_btn = gr.Button(
                    value="Both are bad 👎", 
                    visible=True, 
                    interactive=False, 
                    elem_classes=["vote-btn"],
                )
            with gr.Column():
                refresh_btn = gr.Button(
                    value="New Round 🎲", 
                    interactive=True, 
                    elem_classes=["vote-btn"],
                )
    
    btn_list = [leftvote_btn, rightvote_btn, bothbad_btn, refresh_btn]

    # gen_btn.click(
    #     refresh_side_by_side, 
    #     inputs=[state_models, model_selector_left, model_selector_right], 
    #     outputs=[prompt, input_images, chatbot_left, chatbot_right], 
    #     api_name="refresh_btn_annony"
    # )

    refresh_btn.click(
        refresh_side_by_side, 
        inputs=[state_models, model_selector_left, model_selector_right], 
        outputs=[state0, state1, prompt, input_images, chatbot_left, chatbot_right, leftvote_btn, rightvote_btn, bothbad_btn, model_selector_left, model_selector_right, model_demo_left, model_demo_right], 
        api_name="refresh_btn_annony"
    )

    leftvote_btn.click(
        leftvote_last_response,
        inputs=[state0, state1],
        outputs=[leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right]
    )

    rightvote_btn.click(
        rightvote_last_response,
        inputs=[state0, state1],
        outputs=[leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right]
    )

    bothbad_btn.click(
        bothbad_vote_last_response,
        inputs=[state0, state1],
        outputs=[leftvote_btn, rightvote_btn, bothbad_btn, model_demo_left, model_demo_right]
    )