from data.model_handler import ModelHandler
from app.utils import add_rank_and_format, get_refresh_function
import gradio as gr

METRICS = ["ndcg_at_5", "recall_at_1", "recall_at_5", "mrr_at_5"]

def main(): 
    model_handler = ModelHandler()
    initial_metric = "ndcg_at_5"
    
    data = model_handler.get_vidore_data(initial_metric)
    data = add_rank_and_format(data)

    NUM_DATASETS = len(data.columns) - 3
    NUM_SCORES = len(data) * NUM_DATASETS
    NUM_MODELS = len(data)

    css = """
    table > thead {
        white-space: normal
    }

    table {
        --cell-width-1: 250px
    }

    table > tbody > tr > td:nth-child(2) > div {
        overflow-x: auto
    }

    .filter-checkbox-group {
        max-width: max-content;
    }
    """

    with gr.Blocks(css=css) as block:
        gr.Markdown("# ViDoRe: The Visual Document Retrieval Benchmark 📚🔍")
        gr.Markdown("## From the paper - ColPali: Efficient Document Retrieval with Vision Language Models 👀")

        gr.Markdown(
            """
        Visual Document Retrieval Benchmark leaderboard. To submit, refer to the <a href="https://github.com/tonywu71/vidore-benchmark/" target="_blank" style="text-decoration: underline">ViDoRe GitHub repository</a>.  Refer to the [ColPali paper](https://arxiv.org/abs/XXXX.XXXXX) for details on metrics, tasks and models.
        """
        )
        #all_columns = list(data.columns)
        #default_columns = all_columns

        with gr.Row():
            metric_dropdown = gr.Dropdown(choices=METRICS, value=initial_metric, label="Select Metric")
            #column_checkboxes = gr.CheckboxGroup(choices=all_columns, value=default_columns, label="Select Columns to Display")

        with gr.Row():
            datatype = ["number", "markdown"] + ["number"] * (NUM_DATASETS + 1)
            dataframe = gr.Dataframe(data, datatype=datatype, type="pandas")

        with gr.Row():
            refresh_button = gr.Button("Refresh")
            refresh_button.click(get_refresh_function(), inputs=[metric_dropdown], outputs=dataframe, concurrency_limit=20)


        # Automatically refresh the dataframe when the dropdown value changes
        metric_dropdown.change(get_refresh_function(), inputs=[metric_dropdown], outputs=dataframe)
        #column_checkboxes.change(get_refresh_function(), inputs=[metric_dropdown, column_checkboxes], outputs=dataframe)


        gr.Markdown(
            f"""
        - **Total Datasets**: {NUM_DATASETS}
        - **Total Scores**: {NUM_SCORES}
        - **Total Models**: {NUM_MODELS}
        """
            + r"""
        Please consider citing:

        ```bibtex
        INSERT LATER
        ```
        """
        )

    block.queue(max_size=10).launch(debug=True)


if __name__ == "__main__":    
    main()