import gradio as gr

from src.utils import model_hyperlink

LEADERBOARD_COLUMN_TO_DATATYPE = {
    # open llm
    "Model": "str",
    "Quantization": "str",
    # primary measurements
    "Prefill (tokens/s)": "number",
    "Decode (tokens/s)": "number",
    "Model Size (GB)": "number",
    # deployment settings
    "Backend": "str",
    # additional measurements
    # "Reserved Memory (MB)": "number",
    # "Used Memory (MB)": "number",
    "Params (B)": "number",
    "MMLU Accuracy": "number",
}

PRIMARY_COLUMNS = [
    "Model",
    "Quantization",
    "Prefill (tokens/s)",
    "Decode (tokens/s)",
    "Model Size (GB)",
    "MMLU Accuracy"
]


def process_model(model_name):
    link = f"https://huggingface.co/{model_name}"
    return model_hyperlink(link, model_name)


def get_leaderboard_df(llm_perf_df):
    df = llm_perf_df.copy()
    # transform for leaderboard
    # df["Model"] = df["Model"].apply(process_model)
    return df


def create_leaderboard_table(llm_perf_df):
    # get dataframe
    leaderboard_df = get_leaderboard_df(llm_perf_df)

    # create search bar
    with gr.Row():
        search_bar = gr.Textbox(
            label="Model",
            info="🔍 Search for a model name",
            elem_id="search-bar",
        )
    # create checkboxes
    with gr.Row():
        columns_checkboxes = gr.CheckboxGroup(
            label="Columns 📊",
            value=PRIMARY_COLUMNS,
            choices=list(LEADERBOARD_COLUMN_TO_DATATYPE.keys()),
            info="☑️ Select the columns to display",
            elem_id="columns-checkboxes",
        )
    # create table
    leaderboard_table = gr.components.Dataframe(
        value=leaderboard_df[PRIMARY_COLUMNS],
        datatype=list(LEADERBOARD_COLUMN_TO_DATATYPE.values()),
        headers=list(LEADERBOARD_COLUMN_TO_DATATYPE.keys()),
        elem_id="leaderboard-table",
    )

    return search_bar, columns_checkboxes, leaderboard_table