import gradio as gr

from src.utils import model_hyperlink, process_score


LEADERBOARD_COLUMN_TO_DATATYPE = {
    # open llm
    "Model 🤗": "markdown",
    "Experiment 🧪": "str",
    # primary measurements
    "Prefill Latency (s)": "number",
    "Decode Throughput (tokens/s)": "number",
    "Allocated Memory (MB)": "number",
    "Energy (tokens/kWh)": "number",
    # deployment settings
    "DType 📥": "str",
    "Backend 🏭": "str",
    "Optimization 🛠️": "str",
    "Quantization 🗜️": "str",
    # additional measurements
    "Arch 🏛️": "markdown",
    "Params (B)": "number",
    "Open LLM Score (%)": "number",
    "E2E Latency (s)": "number",
    "E2E Throughput (tokens/s)": "number",
    "Reserved Memory (MB)": "number",
    "Used Memory (MB)": "number",
}


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)
    # process quantization for leaderboard
    df["Open LLM Score (%)"] = df.apply(
        lambda x: process_score(x["Open LLM Score (%)"], x["Quantization 🗜️"]),
        axis=1,
    )
    return df


def create_leaderboard_table(llm_perf_df):
    # get dataframe
    leaderboard_df = get_leaderboard_df(llm_perf_df)
    # create checkboxes
    columns_checkboxes = gr.CheckboxGroup(
        label="Columns 📊",
        choices=list(LEADERBOARD_COLUMN_TO_DATATYPE.keys()),
        value=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,
        datatype=list(LEADERBOARD_COLUMN_TO_DATATYPE.values()),
        headers=list(LEADERBOARD_COLUMN_TO_DATATYPE.keys()),
        elem_id="leaderboard-table",
    )

    return leaderboard_table, columns_checkboxes