import gradio as gr

from src.utils import model_hyperlink, process_score


LEADERBOARD_COLUMN_TO_DATATYPE = {
    # open llm
    "Model 🤗" :"markdown",
    "Arch 🏛️" :"markdown",
    "Params (B)": "number",
    "Open LLM Score (%)": "number",
    # deployment settings
    "DType 📥" :"str",
    "Backend 🏭" :"str",
    "Optimization 🛠️" :"str",
    "Quantization 🗜️" :"str",
    # primary measurements
    "Prefill Latency (s)": "number",
    "Decode Throughput (tokens/s)": "number",
    "Allocated Memory (MB)": "number",
    "Energy (tokens/kWh)": "number",
    # additional measurements
    "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):
    # descriptive text
    gr.HTML("👉 Scroll to the right 👉 for additional columns.", elem_id="text")
    # get dataframe
    leaderboard_df = get_leaderboard_df(llm_perf_df)
    # 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="table",
    )

    return leaderboard_table