import json

import gradio as gr
import pandas as pd
from huggingface_hub import HfFileSystem


RESULTS_DATASET_ID = "datasets/open-llm-leaderboard/results"
EXCLUDED_KEYS =  {
    "pretty_env_info",
    "chat_template",
    "group_subtasks",
}
EXCLUDED_RESULTS_KEYS = {
    "leaderboard",
}
EXCLUDED_RESULTS_LEADERBOARDS_KEYS = {
    "alias",
}

fs = HfFileSystem()


def fetch_result_paths():
    paths = fs.glob(f"{RESULTS_DATASET_ID}/**/**/*.json")
    return paths


def filter_latest_result_path_per_model(paths):
    from collections import defaultdict

    d = defaultdict(list)
    for path in paths:
        model_id, _ = path[len(RESULTS_DATASET_ID) +1:].rsplit("/", 1)
        d[model_id].append(path)
    return {model_id: max(paths) for model_id, paths in d.items()}


def get_result_path_from_model(model_id, result_path_per_model):
    return result_path_per_model[model_id]


def load_data(result_path) -> pd.DataFrame:
    with fs.open(result_path, "r") as f:
        data = json.load(f)
    return data
    # model_name = data.get("model_name", "Model")
    # df = pd.json_normalize([data])
    # return df.iloc[0].rename_axis("Parameters").rename(model_name).to_frame()  # .reset_index()


def load_result(model_id):
    result_path = get_result_path_from_model(model_id, latest_result_path_per_model)
    data = load_data(result_path)
    model_name = data.get("model_name", "Model")
    result = [
        to_vertical(to_dataframe_all(data), model_name),
        to_vertical(to_dataframe_results(data), model_name)
    ]
    return result


def to_dataframe(data):
    return pd.DataFrame.from_records([data])


def to_vertical(df, model_name):
    return df.iloc[0].rename_axis("Parameters").rename(model_name).to_frame()  # .reset_index()


def to_dataframe_all(data):
    return pd.json_normalize([{key: value for key, value in data.items() if key not in EXCLUDED_KEYS}])


def to_dataframe_results(data):
    dfs = {}
    for key in data["results"]:
        if key not in EXCLUDED_RESULTS_KEYS:  # key.startswith("leaderboard_"):
            name = key[len("leaderboard_"):]
            df = to_dataframe(
                {
                    key: value
                    for key, value in data["results"][key].items()
                    if key not in EXCLUDED_RESULTS_LEADERBOARDS_KEYS
                }
            )
            # df.drop(columns=["alias"])
            # df.columns = pd.MultiIndex.from_product([[name], df.columns])
            df.columns = [f"{name}.{column}" for column in df.columns]
            dfs[name] = df
    return pd.concat(dfs.values(), axis="columns")


def concat_result_1(result_1, results):
    return pd.concat([result_1, results.iloc[:, [0, 2]].set_index("Parameters")], axis=1).reset_index()


def concat_result_2(result_2, results):
    return pd.concat([results.iloc[:, [0, 1]].set_index("Parameters"), result_2], axis=1).reset_index()


def render_result_1(model_id, *results):
    result = load_result(model_id)
    return [concat_result_1(*result_args) for result_args in zip(result, results)]


def render_result_2(model_id, *results):
    result = load_result(model_id)
    return [concat_result_2(*result_args) for result_args in zip(result, results)]


# if __name__ == "__main__":
latest_result_path_per_model = filter_latest_result_path_per_model(fetch_result_paths())

with gr.Blocks(fill_height=True) as demo:
    gr.HTML("<h1 style='text-align: center;'>Compare Results of the 🤗 Open LLM Leaderboard</h1>")
    gr.HTML("<h3 style='text-align: center;'>Select 2 results to load and compare</h3>")

    with gr.Row():
        with gr.Column():
            model_id_1 = gr.Dropdown(choices=list(latest_result_path_per_model.keys()), label="Results")
            load_btn_1 = gr.Button("Load")
        with gr.Column():
            model_id_2 = gr.Dropdown(choices=list(latest_result_path_per_model.keys()), label="Results")
            load_btn_2 = gr.Button("Load")

    with gr.Row():
        with gr.Tab("All"):
            compared_results_all = gr.Dataframe(
                label="Results",
                headers=["Parameters", "Model-1", "Model-2"],
                interactive=False,
                column_widths=["30%", "30%", "30%"],
                wrap=True,
            )
        with gr.Tab("Results"):
            compared_results_results = gr.Dataframe(
                label="Results",
                headers=["Parameters", "Model-1", "Model-2"],
                interactive=False,
                column_widths=["30%", "30%", "30%"],
                wrap=True,
            )

    load_btn_1.click(
        fn=render_result_1,
        inputs=[model_id_1, compared_results_all, compared_results_results],
        outputs=[compared_results_all, compared_results_results],
    )
    load_btn_2.click(
        fn=render_result_2,
        inputs=[model_id_2, compared_results_all, compared_results_results],
        outputs=[compared_results_all, compared_results_results],
    )

demo.launch()