import gradio as gr
from lb_info import load_results, BUILD_L1_DF
from desc import (
    LEADERBOARD_INTRODUCTION,
    LEADERBOARD_MD,
    CITATION_BUTTON_TEXT,
    CITATION_BUTTON_LABEL,
    CINEPILE_ABOUT_MD,
)
from urllib.request import urlopen
import markdown


def filter_df(fields):
    # Use set operations to avoid duplicates
    headers = (
        [
            "Model",
            "Params (B)",
            "Average Accuracy",
        ]
        + fields
        + [
            "Average Rank",
        ]
    )

    # Remove duplicates in headers by keeping the earliest entry
    headers = list(dict.fromkeys(headers))
    return table[headers]


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    struct = load_results()
    results = struct

    # Build leaderboard DataFrame for CinePile data
    table, check_box = BUILD_L1_DF(results)

    N_MODELS = len(table)
    UP_TS = "20th October 2024"  # Replace with actual timestamp

    gr.Markdown(LEADERBOARD_INTRODUCTION.format(N_MODELS, UP_TS))

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        # First Tab: CinePile Leaderboard
        with gr.TabItem("CinePile Leaderboard", elem_id="main"):
            gr.Markdown(LEADERBOARD_MD)

            # Checkbox for selecting question categories
            checkbox_group = gr.CheckboxGroup(
                choices=check_box["question_categories"],
                label="Question Categories",
                interactive=True,
            )

            # DataFrame component for displaying the leaderboard
            data_component = gr.DataFrame(
                value=table[check_box["essential"]],
                datatype=[check_box["type_map"][x] for x in check_box["essential"]],
                interactive=False,
                visible=True,
            )

            # Update the table when checkbox changes
            checkbox_group.change(
                fn=filter_df, inputs=checkbox_group, outputs=data_component
            )

        # Second Tab: About
        with gr.TabItem("About CinePile", elem_id="about"):
            markdown_content = urlopen(CINEPILE_ABOUT_MD).read().decode()
            gr.HTML(markdown.markdown(markdown_content))

        # Add citation support under "About"
        with gr.Row():
            with gr.Accordion("Citation", open=False):
                citation_button = gr.Textbox(
                    value=CITATION_BUTTON_TEXT,
                    label=CITATION_BUTTON_LABEL,
                    elem_id="citation-button",
                )

demo.launch()