from pathlib import Path

from apscheduler.schedulers.background import BackgroundScheduler
import pandas as pd
import gradio as gr
from gradio_leaderboard import Leaderboard, ColumnFilter

from constants import Constants, model_type_emoji


TITLE = """<h1 align="center" id="space-title">TabArena: Public leaderboard for Tabular methods</h1>"""

INTRODUCTION_TEXT = ("TabArena Leaderboard measures the performance of tabular models on a collection of tabular "
                     "datasets manually curated. The datasets are collected to make sure they are tabular, with "
                     "permissive license without ethical issues and so on, we refer to the paper for a full "
                     "description of our approach.")

ABOUT_TEXT = f"""
## How It Works.

To evaluate the leaderboard, follow install instructions in 
`https://github.com/autogluon/tabrepo/tree/tabarena` and run 
`https://github.com/autogluon/tabrepo/blob/tabarena/examples/tabarena/run_tabarena_eval.py`.


This will generate a leaderboard. You can add your own method and contact the authors if you want it to be added
to the leaderboard. We require method to have public code available to be considered in the leaderboard. 
"""

CITATION_BUTTON_LABEL = "If you use this leaderboard in your research please cite the following:"
CITATION_BUTTON_TEXT = r"""
@article{
TBA,
}
"""


def get_model_family(model_name: str) -> str:
    prefixes_mapping = {
        Constants.automl: ["AutoGluon"],
        Constants.finetuned: ["REALMLP", "TabM", "FASTAI", "MNCA", "NN_TORCH"],
        Constants.tree: ["GBM", "CAT", "EBM", "XGB"],
        Constants.foundational: ["TABDPT", "TABICL", "TABPFN"],
        Constants.baseline: ["KNN", "LR"]
    }
    for method_type, prefixes in prefixes_mapping.items():
        for prefix in prefixes:
            if prefix.lower() in model_name.lower():
                return method_type
    return Constants.other


def load_data(filename: str):
    df_leaderboard = pd.read_csv(Path(__file__).parent / "data" / f"{filename}.csv.zip")
    print(f"Loaded dataframe with {len(df_leaderboard)} rows and columns {df_leaderboard.columns}")

    # sort by ELO
    df_leaderboard.sort_values(by="elo", ascending=False, inplace=True)

    # add model family information
    df_leaderboard["family"] = df_leaderboard.loc[:, "method"].apply(
        lambda s: get_model_family(s) + " " + model_type_emoji[get_model_family(s)]
    )

    # select only the columns we want to display
    df_leaderboard = df_leaderboard.loc[:, ["method", "family", "time_train_s", "time_infer_s", "rank", "elo"]]

    # round for better display
    df_leaderboard = df_leaderboard.round(1)

    # rename some columns
    df_leaderboard.rename(columns={
        "time_train_s": "training time (s)",
        "time_infer_s": "inference time (s)",
    }, inplace=True)

    # TODO show ELO +/- sem
    return df_leaderboard


def make_leaderboard(df_leaderboard: pd.DataFrame) -> Leaderboard:
    return Leaderboard(
        value=df_leaderboard,
        search_columns=["method"],
        filter_columns=[
            # "method",
            ColumnFilter("family", type="dropdown", label="Filter by family"),
        ]
    )


def main():

    demo = gr.Blocks()
    with demo:
        gr.HTML(TITLE)
        gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")

        with gr.Tabs(elem_classes="tab-buttons") as tabs:
            with gr.TabItem('🏅 Overall', elem_id="llm-benchmark-tab-table", id=2):
                df_leaderboard = load_data("leaderboard-all")
                leaderboard = make_leaderboard(df_leaderboard)

            with gr.TabItem('🏅 Regression', elem_id="llm-benchmark-tab-table", id=0):
                df_leaderboard = load_data("leaderboard-regression")
                leaderboard = make_leaderboard(df_leaderboard)

            with gr.TabItem('🏅 Classification', elem_id="llm-benchmark-tab-table", id=1):
                df_leaderboard = load_data("leaderboard-classification")
                leaderboard = make_leaderboard(df_leaderboard)

            with gr.TabItem("📝 About", elem_id="llm-benchmark-tab-table", id=4):
                gr.Markdown(ABOUT_TEXT, elem_classes="markdown-text")
        with gr.Row():
            with gr.Accordion("📙 Citation", open=False):
                citation_button = gr.Textbox(
                    value=CITATION_BUTTON_TEXT,
                    label=CITATION_BUTTON_LABEL,
                    lines=20,
                    elem_id="citation-button",
                    show_copy_button=True,
                )

    scheduler = BackgroundScheduler()
    # scheduler.add_job(restart_space, "interval", seconds=1800)
    scheduler.start()
    demo.queue(default_concurrency_limit=40).launch()
    demo.launch()


if __name__ == "__main__":
    main()