import gradio as gr

import src.dependency  # noqa
from src.assets import custom_css

# from src.attention import create_attn_plots
from src.content import ABOUT, CITATION_BUTTON, CITATION_BUTTON_LABEL, LOGO, TITLE
from src.hardware import load_hardware_configs
from src.leaderboard import create_leaderboard_table
from src.llm_perf import get_llm_perf_df
from src.map import create_lat_score_mem_plot
from src.panel import (
    create_control_callback,
    create_control_panel,
    create_select_callback,
)

configs = load_hardware_configs("hardware.yaml")


demo = gr.Blocks(
    css=custom_css,
    theme=gr.themes.Default(primary_hue="indigo", secondary_hue="indigo"),
)
with demo:
    print("Building demo")
    gr.HTML(LOGO, elem_classes="logo")
    gr.HTML(TITLE, elem_classes="title")
    ####################### HARDWARE TABS #######################
    with gr.Tabs(elem_classes="tabs"):
        for id, config in enumerate(configs):
            with gr.TabItem(config.description, id=id):
                ####################### HARDWARE DETAILS #######################
                if config.detail:
                    gr.Markdown(config.detail, elem_classes="descriptive-text")

                # ####################### CONTROL PANEL #######################
                (
                    filter_button,
                    machine_value,
                    subsets_value,
                    backends_value,
                    hardware_type_value,
                    score_slider,
                    memory_slider,
                    backend_checkboxes,
                    datatype_checkboxes,
                    optimization_checkboxes,
                    quantization_checkboxes,
                    kernels_checkboxes,
                ) = create_control_panel(
                    machine=config.machine,
                    subsets=config.subsets,
                    backends=config.backends,
                    hardware_type=config.hardware_type,
                    hardware_provider=config.hardware_provider,
                )
                ####################### HARDWARE SUBTABS #######################
                with gr.Tabs(elem_classes="subtabs"):
                    open_llm_perf_df = get_llm_perf_df(
                        machine=config.machine,
                        subsets=config.subsets,
                        backends=config.backends,
                        hardware_type=config.hardware_type,
                    )
                    ####################### LEADERBOARD TAB #######################
                    with gr.TabItem("Leaderboard 🏅", id=0):
                        search_bar, columns_checkboxes, leaderboard_table = (
                            create_leaderboard_table(open_llm_perf_df)
                        )
                #     with gr.TabItem("Find Your Best Model 🧭", id=1):
                #         lat_score_mem_plot = create_lat_score_mem_plot(
                #             open_llm_perf_df
                #         )


                # ####################### CONTROL CALLBACK #######################
                # create_control_callback(
                #     filter_button,
                #     # inputs
                #     machine_value,
                #     subsets_value,
                #     backends_value,
                #     hardware_type_value,
                #     score_slider,
                #     memory_slider,
                #     backend_checkboxes,
                #     datatype_checkboxes,
                #     optimization_checkboxes,
                #     quantization_checkboxes,
                #     kernels_checkboxes,
                #     # interactive
                #     columns_checkboxes,
                #     search_bar,
                #     # outputs
                #     leaderboard_table,
                #     lat_score_mem_plot,
                #     # attn_prefill_plot,
                #     # attn_decode_plot,
                #     # quant_krnl_prefill_plot,
                #     # quant_krnl_decode_plot,
                # )

                # create_select_callback(
                #     # inputs
                #     machine_value,
                #     subsets_value,
                #     backends_value,
                #     hardware_type_value,
                #     # interactive
                #     columns_checkboxes,
                #     search_bar,
                #     # outputs
                #     leaderboard_table,
                # )

    ####################### ABOUT TAB #######################
    with gr.TabItem("About 📖", id=len(configs)):
        gr.Markdown(ABOUT, elem_classes="descriptive-text")
    ####################### CITATION
    with gr.Row():
        with gr.Accordion("📙 Citation", open=False):
            citation_button = gr.Textbox(
                value=CITATION_BUTTON,
                label=CITATION_BUTTON_LABEL,
                elem_id="citation-button",
                show_copy_button=True,
            )
    print("Building demo end")

if __name__ == "__main__":
    demo.queue().launch(ssr_mode=False)