__all__ = ['block']

import gradio as gr
import pandas as pd
from constants import (
    MODEL_INFO, TASK_INFO, AVG_INFO, DATA_TITILE_TYPE,
    COLUMN_NAMES, CSV_RESULT_PATH, LEADERBORAD_INTRODUCTION,
    CITATION_BUTTON_TEXT, CITATION_BUTTON_LABEL
)


def get_baseline_df(selected_columns=None):
    if selected_columns is None:
        selected_columns = AVG_INFO
    df = pd.read_csv(CSV_RESULT_PATH)
    df = df.sort_values(by="Overall", ascending=False)
    present_columns = MODEL_INFO + selected_columns
    df = df[present_columns]
    return df


def get_all_df():
    df = pd.read_csv(CSV_RESULT_PATH)
    df = df.sort_values(by="Overall", ascending=False)
    return df


block = gr.Blocks()


with block:
    gr.Markdown(
        LEADERBORAD_INTRODUCTION
    )
    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        # table jmmmu bench
        with gr.TabItem("🏅 JMMMU-Pro Benchmark", elem_id="jmmmu-pro-benchmark-tab-table", id=1):
            # selection for column part:
            checkbox_group = gr.CheckboxGroup(
                choices=TASK_INFO,
                value=AVG_INFO,
                label="Evaluation Dimension",
                interactive=True,
            )  # user can select the evaluation dimension

            baseline_value = get_baseline_df(checkbox_group.value)
            baseline_header = MODEL_INFO + checkbox_group.value
            baseline_datatype = ['markdown'] * 2 + ['number'] * len(checkbox_group.value)

            data_component = gr.components.Dataframe(
                value=baseline_value,
                headers=baseline_header,
                type="pandas",
                datatype=baseline_datatype,
                interactive=False,
                visible=True,
                )

            def on_filter_method_change(selected_columns):
                updated_data = get_all_df()

                # columns:
                selected_columns = [
                    item for item in TASK_INFO if item in selected_columns
                ]
                present_columns = MODEL_INFO + selected_columns
                updated_data = updated_data[present_columns]
                updated_data = updated_data.sort_values(
                    by=selected_columns[0], ascending=False
                )
                updated_headers = present_columns
                update_datatype = [
                    DATA_TITILE_TYPE[COLUMN_NAMES.index(x)]
                    for x in updated_headers
                ]

                filter_component = gr.components.Dataframe(
                    value=updated_data,
                    headers=updated_headers,
                    type="pandas",
                    datatype=update_datatype,
                    interactive=False,
                    visible=True,
                    )
                return filter_component

            checkbox_group.change(fn=on_filter_method_change, inputs=[checkbox_group], outputs=data_component)

    def refresh_data():
        value = get_baseline_df(checkbox_group.value)
        return value

    with gr.Row():
        data_run = gr.Button("Refresh")
        data_run.click(
            refresh_data, outputs=[data_component]
        )

    with gr.Accordion("Citation", open=False):
        citation_button = gr.Textbox(
            value=CITATION_BUTTON_TEXT,
            label=CITATION_BUTTON_LABEL,
            elem_id="citation-button",
            show_copy_button=True,
        )

block.launch()