import gradio as gr

from src.utils import model_hyperlink, process_score


LEADERBOARD_COLUMN_TO_DATATYPE = {
    # open llm
    "Model 🤗" :"str",
    "Arch 🏛️" :"str",
    "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",
}

from dataclasses import dataclass

@dataclass
class LeaderboardColumn:
    name: str
    type: str 

LEADERBOARD_COLUMNS = [
    LeaderboardColumn("Model 🤗", "str"),
    LeaderboardColumn("Arch 🏛️", "str"),
    LeaderboardColumn("Params (B)", "number"),
    LeaderboardColumn("Open LLM Score (%)", "number"),
    LeaderboardColumn("DType 📥", "str"),
    LeaderboardColumn("Backend 🏭", "str"),
    LeaderboardColumn("Optimization 🛠️", "str"),
    LeaderboardColumn("Quantization 🗜️", "str"),
    LeaderboardColumn("Prefill Latency (s)", "number"),
    LeaderboardColumn("Decode Throughput (tokens/s)", "number"),
    LeaderboardColumn("Allocated Memory (MB)", "number"),
    LeaderboardColumn("Energy (tokens/kWh)", "number"),
    LeaderboardColumn("E2E Latency (s)", "number"),
    LeaderboardColumn("E2E Throughput (tokens/s)", "number"),
    LeaderboardColumn("Reserved Memory (MB)", "number"),
    LeaderboardColumn( "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

COLS = [col.name for col in LEADERBOARD_COLUMNS]
TYPES = [col.type for col in LEADERBOARD_COLUMNS]

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

    return leaderboard_df