import os
import gradio as gr
import pandas as pd
import plotly.express as px
from apscheduler.schedulers.background import BackgroundScheduler

from src.assets.css_html_js import custom_css, custom_js
from src.assets.text_content import (
    TITLE,
    INTRODUCTION_TEXT,
    ABOUT_TEXT,
    EXAMPLE_CONFIG_TEXT,
    CITATION_BUTTON_LABEL,
    CITATION_BUTTON_TEXT,
)
from src.utils import (
    change_tab,
    restart_space,
    load_dataset_repo,
    process_model_name,
    process_model_type,
)


LLM_PERF_LEADERBOARD_REPO = "optimum/llm-perf-leaderboard"
LLM_PERF_DATASET_REPO = "optimum/llm-perf-dataset"
OPTIMUM_TOKEN = os.environ.get("OPTIMUM_TOKEN", None)

ALL_COLUMNS_MAPPING = {
    "weight_class": "Weight Class 🏋️",
    "model_type": "LLM Type 🤗",
    "best_scored_model": "Best Scored LLM 🏆",
    #
    "backend.name": "Backend 🏭",
    "backend.torch_dtype": "Dtype 📥",
    "quantization": "Quantization 🗜️",
    "optimizations": "Optimizations 🛠️",
    #
    "best_score": "Best Score (%) ⬆️",
    "generate.peak_memory(MB)": "Memory (MB) ⬇️",
    "generate.throughput(tokens/s)": "Throughput (tokens/s) ⬆️",
    "generate.energy_consumption(kWh/token)": "Energy (kWh/token) ⬇️",
    #
}
ALL_COLUMNS_DATATYPES = [
    "str",
    "str",
    "markdown",
    #
    "str",
    "str",
    "str",
    "str",
    #
    "str",
    "number",
    "number",
    "number",
    #
]
SORTING_COLUMN = ["perf_distance"]

llm_perf_dataset_repo = load_dataset_repo(LLM_PERF_DATASET_REPO, OPTIMUM_TOKEN)


def get_benchmark_df(benchmark="Succeeded-1xA100-80GB"):
    if llm_perf_dataset_repo:
        llm_perf_dataset_repo.git_pull()

    # load data
    benchmark_df = pd.read_csv(f"./llm-perf-dataset/reports/{benchmark}.csv")
    clusters_df = pd.read_csv("./llm-perf-dataset/Clustered-Open-LLM-Leaderboard.csv")
    # merge on model
    merged_df = benchmark_df.merge(
        clusters_df, left_on="model", right_on="best_scored_model"
    )
    # fix energy consumption nans
    merged_df["generate.energy_consumption(kWh/token)"].fillna("N/A", inplace=True)

    # add optimizations
    merged_df["optimizations"] = merged_df["backend.bettertransformer"].apply(
        lambda x: "BetterTransformer" if x else "None"
    )
    # add quantization scheme
    merged_df["quantization"] = merged_df["backend.quantization_strategy"].apply(
        lambda x: "BnB.4bit" if x == "bnb" else ("GPTQ.4bit" if x == "gptq" else "None")
    )
    # distance to 100% score, normalized to 0, 1
    score_distance = (100 - merged_df["best_score"]) / 100
    # distance to 0s latency, normalized to 0, 1
    latency_distance = merged_df["generate.latency(s)"] / (
        merged_df["generate.latency(s)"].max() - merged_df["generate.latency(s)"].min()
    )
    # distance to 0MB memory
    memory_distance = merged_df["forward.peak_memory(MB)"] / (
        merged_df["forward.peak_memory(MB)"].max()
        - merged_df["forward.peak_memory(MB)"].min()
    )
    # add perf distance
    merged_df["perf_distance"] = (
        score_distance**2 + latency_distance**2 + memory_distance**2
    ) ** 0.5

    return merged_df


def get_benchmark_table(bench_df):
    # add * to quantized models score
    copy_df = bench_df.copy()
    # add * to quantized models score since we can't garantee the score is the same
    copy_df["best_score"] = copy_df.apply(
        lambda x: f"{x['best_score']}**"
        if x["backend.quantization_strategy"]
        else x["best_score"],
        axis=1,
    )
    # sort
    copy_df.sort_values(by=SORTING_COLUMN, ascending=True, inplace=True)
    # filter
    copy_df = copy_df[list(ALL_COLUMNS_MAPPING.keys())]
    # rename
    copy_df.rename(columns=ALL_COLUMNS_MAPPING, inplace=True)
    # transform
    copy_df["LLM Type 🤗"] = copy_df["LLM Type 🤗"].apply(process_model_type)
    copy_df["Best Scored LLM 🏆"] = copy_df["Best Scored LLM 🏆"].apply(
        process_model_name
    )
    return copy_df


def get_benchmark_plot(bench_df):
    fig = px.scatter(
        bench_df,
        y="best_score",
        x="generate.throughput(tokens/s)",
        size="generate.peak_memory(MB)",
        color="model_type",
        custom_data=list(ALL_COLUMNS_MAPPING.keys()),
        color_discrete_sequence=px.colors.qualitative.Light24,
    )
    fig.update_layout(
        title={
            "text": "Model Score vs. Latency vs. Memory",
            "y": 0.95,
            "x": 0.5,
            "xanchor": "center",
            "yanchor": "top",
        },
        xaxis_title="Generation Throughput (tokens/s)",
        yaxis_title="Open LLM Score (%)",
        legend_title="Model Type",
        width=1200,
        height=600,
    )
    fig.update_traces(
        hovertemplate="<br>".join(
            [
                f"<b>{ALL_COLUMNS_MAPPING[key]}:</b> %{{customdata[{i}]}}"
                for i, key in enumerate(ALL_COLUMNS_MAPPING.keys())
            ]
        )
    )
    return fig


def filter_query(
    text,
    backends,
    datatypes,
    optimizations,
    score,
    memory,
    benchmark="Succeeded-1xA100-80GB",
):
    raw_df = get_benchmark_df(benchmark=benchmark)
    filtered_df = raw_df[
        raw_df["best_scored_model"].str.lower().str.contains(text.lower())
        & raw_df["backend.name"].isin(backends)
        & raw_df["backend.torch_dtype"].isin(datatypes)
        & (
            pd.concat(
                [
                    raw_df["optimizations"].str.contains(optimization)
                    for optimization in optimizations
                ],
                axis=1,
            ).any(axis="columns")
            if len(optimizations) > 0
            else True
        )
        & (raw_df["best_score"] >= score)
        & (raw_df["forward.peak_memory(MB)"] <= memory)
    ]
    filtered_table = get_benchmark_table(filtered_df)
    filtered_plot = get_benchmark_plot(filtered_df)
    return filtered_table, filtered_plot


# Dataframes
A100_df = get_benchmark_df(benchmark="Succeeded-1xA100-80GB")
A100_table = get_benchmark_table(A100_df)
A100_plot = get_benchmark_plot(A100_df)

# Demo interface
demo = gr.Blocks(css=custom_css)
with demo:
    # leaderboard title
    gr.HTML(TITLE)
    # introduction text
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="descriptive-text")

    # leaderboard tabs
    with gr.Tabs(elem_classes="A100-tabs") as A100_tabs:
        with gr.TabItem("🖥️ A100-80GB Benchmark 🏆", id=0):
            gr.HTML(
                "👉 Scroll to the right 👉 for more columns.", elem_id="descriptive-text"
            )
            # Original leaderboard table
            A100_leaderboard = gr.components.Dataframe(
                value=A100_table,
                datatype=ALL_COLUMNS_DATATYPES,
                headers=list(ALL_COLUMNS_MAPPING.values()),
                elem_id="1xA100-table",
            )

        with gr.TabItem("🖥️ A100-80GB Plot 📊", id=1):
            gr.HTML(
                "👆 Hover over the points 👆 for additional information.",
                elem_id="descriptive-text",
            )
            # Original leaderboard plot
            A100_plotly = gr.components.Plot(
                value=A100_plot,
                elem_id="1xA100-plot",
                show_label=False,
            )

        with gr.TabItem("Control Panel 🎛️", id=2):
            gr.HTML(
                "Use this control panel to filter the leaderboard's table and plot.",
                elem_id="descriptive-text",
            )
            # control panel interface
            with gr.Row():
                with gr.Column(scale=1):
                    search_bar = gr.Textbox(
                        label="Model 🤗",
                        info="🔍 Search for a model name",
                        elem_id="search-bar",
                    )
                with gr.Column(scale=1):
                    with gr.Box():
                        score_slider = gr.Slider(
                            label="Open LLM Score 📈",
                            info="🎚️ Slide to minimum Open LLM score",
                            value=0,
                            elem_id="threshold-slider",
                        )
                with gr.Column(scale=1):
                    with gr.Box():
                        memory_slider = gr.Slider(
                            label="Peak Memory (MB) 📈",
                            info="🎚️ Slide to maximum Peak Memory",
                            minimum=0,
                            maximum=80 * 1024,
                            value=80 * 1024,
                            elem_id="memory-slider",
                        )

            with gr.Row():
                with gr.Column(scale=1):
                    backend_checkboxes = gr.CheckboxGroup(
                        label="Backends 🏭",
                        choices=["pytorch", "onnxruntime"],
                        value=["pytorch", "onnxruntime"],
                        info="☑️ Select the backends",
                        elem_id="backend-checkboxes",
                    )
                with gr.Column(scale=1):
                    datatype_checkboxes = gr.CheckboxGroup(
                        label="Dtypes 📥",
                        choices=["float32", "float16"],
                        value=["float32", "float16"],
                        info="☑️ Select the load dtypes",
                        elem_id="dtype-checkboxes",
                    )
                with gr.Column(scale=2):
                    optimizations_checkboxes = gr.CheckboxGroup(
                        label="Optimizations 🛠️",
                        choices=["None", "BetterTransformer"],
                        value=["None", "BetterTransformer"],
                        info="☑️ Select the optimizations",
                        elem_id="optimizations-checkboxes",
                    )

            with gr.Row():
                filter_button = gr.Button(
                    value="Filter 🚀",
                    elem_id="filter-button",
                )

        with gr.TabItem("About 📖", id=3):
            gr.HTML(ABOUT_TEXT, elem_classes="descriptive-text")
            gr.Markdown(EXAMPLE_CONFIG_TEXT, elem_classes="descriptive-text")

    demo.load(
        change_tab,
        A100_tabs,
        _js=custom_js,
    )

    filter_button.click(
        filter_query,
        [
            search_bar,
            backend_checkboxes,
            datatype_checkboxes,
            optimizations_checkboxes,
            score_slider,
            memory_slider,
        ],
        [A100_leaderboard, A100_plotly],
    )

    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",
            ).style(show_copy_button=True)


# Restart space every hour
scheduler = BackgroundScheduler()
scheduler.add_job(
    restart_space,
    "interval",
    seconds=3600,
    args=[LLM_PERF_LEADERBOARD_REPO, OPTIMUM_TOKEN],
)
scheduler.start()

# Launch demo
demo.queue(concurrency_count=40).launch()