import json

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

from src.constants import DETAILS_DATASET_ID, DETAILS_FILENAME, SUBTASKS, TASKS
from src.results import fetch_result_paths, filter_latest_result_path_per_model, update_load_results_component, \
    load_results_dataframes, display_results, update_tasks_component, clear_results

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 update_load_results_component():
    return gr.Button("Load Results", interactive=True)


def load_data(result_path) -> pd.DataFrame:
    with fs.open(result_path, "r") as f:
        data = json.load(f)
    return data


def load_results_dataframe(model_id):
    if not model_id:
        return
    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")
    df = pd.json_normalize([{key: value for key, value in data.items()}])
    # df.columns = df.columns.str.split(".")  # .split return a list instead of a tuple
    return df.set_index(pd.Index([model_name])).reset_index()


def load_results_dataframes(*model_ids):
    return [load_results_dataframe(model_id) for model_id in model_ids]


def display_results(task, *dfs):
    dfs = [df.set_index("index") for df in dfs if "index" in df.columns]
    if not dfs:
        return None, None
    df = pd.concat(dfs)
    df = df.T.rename_axis(columns=None)
    return display_tab("results", df, task), display_tab("configs", df, task)


def display_tab(tab, df, task):
    df = df.style.format(na_rep="")
    df.hide(
        [
            row
            for row in df.index
            if (
                not row.startswith(f"{tab}.")
                or row.startswith(f"{tab}.leaderboard.")
                or row.endswith(".alias")
                or (not row.startswith(f"{tab}.{task}") if task != "All" else False)
            )
        ],
        axis="index",
    )
    start = len(f"{tab}.leaderboard_") if task == "All" else len(f"{tab}.{task} ")
    df.format_index(lambda idx: idx[start:].removesuffix(",none"), axis="index")
    return df.to_html()


def update_tasks_component():
    return gr.Radio(
            ["All"] + list(TASKS.values()),
            label="Tasks",
            info="Evaluation tasks to be displayed",
            value="All",
            interactive=True,
        )


def clear_results():
    # model_id_1, model_id_2, dataframe_1, dataframe_2, task
    return (
        None, None, None, None,
        gr.Radio(
            ["All"] + list(TASKS.values()),
            label="Tasks",
            info="Evaluation tasks to be displayed",
            value="All",
            interactive=False,
        ),
    )


def update_subtasks_component(task):
    return gr.Radio(
        SUBTASKS.get(task),
        info="Evaluation subtasks to be displayed",
        value=None,
    )


def update_load_details_component(model_id_1, model_id_2, subtask):
    if (model_id_1 or model_id_2) and subtask:
        return gr.Button("Load Details", interactive=True)
    else:
        return gr.Button("Load Details", interactive=False)


def load_details_dataframe(model_id, subtask):
    if not model_id or not subtask:
        return
    model_name_sanitized = model_id.replace("/", "__")
    paths = fs.glob(
        f"{DETAILS_DATASET_ID}/**/{DETAILS_FILENAME}".format(
            model_name_sanitized=model_name_sanitized, subtask=subtask
        )
    )
    if not paths:
        return
    path = max(paths)
    with fs.open(path, "r") as f:
        data = [json.loads(line) for line in f]
    df = pd.json_normalize(data)
    # df = df.rename_axis("Parameters", axis="columns")
    df["model_name"] = model_id  # Keep model_name
    return df
    # return df.set_index(pd.Index([model_id])).reset_index()


def load_details_dataframes(subtask, *model_ids):
    return [load_details_dataframe(model_id, subtask) for model_id in model_ids]


def display_details(sample_idx, *dfs):
    rows = [df.iloc[sample_idx] for df in dfs if "model_name" in df.columns and sample_idx < len(df)]
    if not rows:
        return
    # Pop model_name and add it to the column name
    df = pd.concat([row.rename(row.pop("model_name")) for row in rows], axis="columns")
    return (
        df.style
        .format(na_rep="")
        # .hide(axis="index")
        .to_html()
    )


def update_sample_idx_component(*dfs):
    maximum = max([len(df) - 1 for df in dfs])
    return gr.Number(
        label="Sample Index",
        info="Index of the sample to be displayed",
        value=0,
        minimum=0,
        maximum=maximum,
        visible=True,
    )


def clear_details():
    # model_id_1, model_id_2, details_dataframe_1, details_dataframe_2, details_task, subtask, sample_idx
    return (
        None, None, None, None, None, None,
        gr.Number(label="Sample Index", info="Index of the sample to be displayed", value=0, minimum=0,visible=False),
    )


# 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 models to load and compare their results</h3>")

    with gr.Row():
        with gr.Column():
            model_id_1 = gr.Dropdown(choices=list(latest_result_path_per_model.keys()), label="Models")
            dataframe_1 = gr.Dataframe(visible=False)
        with gr.Column():
            model_id_2 = gr.Dropdown(choices=list(latest_result_path_per_model.keys()), label="Models")
            dataframe_2 = gr.Dataframe(visible=False)

    with gr.Row():
        # with gr.Tab("All"):
        #     pass
        with gr.Tab("Results"):
            task = gr.Radio(
                ["All"] + list(TASKS.values()),
                label="Tasks",
                info="Evaluation tasks to be displayed",
                value="All",
                interactive=False,
            )
            load_results_btn = gr.Button("Load Results", interactive=False)
            clear_results_btn = gr.Button("Clear Results")
            with gr.Tab("Results"):
                results = gr.HTML()
            with gr.Tab("Configs"):
                configs = gr.HTML()
        with gr.Tab("Details"):
            details_task = gr.Radio(
                ["All"] + list(TASKS.values()),
                label="Tasks",
                info="Evaluation tasks to be displayed",
                value="All",
                interactive=True,
            )
            subtask = gr.Radio(
                SUBTASKS.get(details_task.value),
                label="Subtasks",
                info="Evaluation subtasks to be displayed (choose one of the Tasks above)",
            )
            load_details_btn = gr.Button("Load Details", interactive=False)
            clear_details_btn = gr.Button("Clear Details")
            sample_idx = gr.Number(
                label="Sample Index",
                info="Index of the sample to be displayed",
                value=0,
                minimum=0,
                visible=False
            )
            details = gr.HTML()
            details_dataframe_1 = gr.Dataframe(visible=False)
            details_dataframe_2 = gr.Dataframe(visible=False)
            details_dataframe = gr.DataFrame(visible=False)

    model_id_1.change(
        fn=update_load_results_component,
        outputs=load_results_btn,
    )
    load_results_btn.click(
        fn=load_results_dataframes,
        inputs=[model_id_1, model_id_2],
        outputs=[dataframe_1, dataframe_2],
    ).then(
        fn=update_tasks_component,
        outputs=task,
    )
    gr.on(
        triggers=[dataframe_1.change, dataframe_2.change, task.change],
        fn=display_results,
        inputs=[task, dataframe_1, dataframe_2],
        outputs=[results, configs],
    )
    clear_results_btn.click(
        fn=clear_results,
        outputs=[model_id_1, model_id_2, dataframe_1, dataframe_2, task],
    )

    details_task.change(
        fn=update_subtasks_component,
        inputs=details_task,
        outputs=subtask,
    )
    gr.on(
        triggers=[model_id_1.change, model_id_2.change, subtask.change, details_task.change],
        fn=update_load_details_component,
        inputs=[model_id_1, model_id_2, subtask],
        outputs=load_details_btn,
    )
    load_details_btn.click(
        fn=load_details_dataframes,
        inputs=[subtask, model_id_1, model_id_2],
        outputs=[details_dataframe_1, details_dataframe_2],
    ).then(
        fn=update_sample_idx_component,
        inputs=[details_dataframe_1, details_dataframe_2],
        outputs=sample_idx,
    )
    gr.on(
        triggers=[details_dataframe_1.change, details_dataframe_2.change, sample_idx.change],
        fn=display_details,
        inputs=[sample_idx, details_dataframe_1, details_dataframe_2],
        outputs=details,
    )
    clear_details_btn.click(
        fn=clear_details,
        outputs=[model_id_1, model_id_2, details_dataframe_1, details_dataframe_2, details_task, subtask, sample_idx],
    )

demo.launch()