import json

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


RESULTS_DATASET_ID = "datasets/open-llm-leaderboard/results"


fs = HfFileSystem()


def fetch_results():
    files = fs.glob(f"{RESULTS_DATASET_ID}/**/**/*.json")
    results = [file[len(RESULTS_DATASET_ID) +1:] for file in files]
    return results


def load_result(result_path) -> pd.DataFrame:
    with fs.open(f"{RESULTS_DATASET_ID}/{result_path}", "r") as f:
        data = json.load(f)
    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 render_result_1(result_path, results):
    result = load_result(result_path)
    return pd.concat([result, results.iloc[:, [0, 2]].set_index("Parameters")], axis=1).reset_index()


def render_result_2(result_path, results):
    result = load_result(result_path)
    return pd.concat([results.iloc[:, [0, 1]].set_index("Parameters"), result], axis=1).reset_index()


if __name__ == "__main__":
    results = fetch_results()

    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():
                result_path_1 = gr.Dropdown(choices=results, label="Results")
                load_btn_1 = gr.Button("Load")
            with gr.Column():
                result_path_2 = gr.Dropdown(choices=results, label="Results")
                load_btn_2 = gr.Button("Load")

        with gr.Row():
            compared_results = gr.Dataframe(
                label="Results",
                headers=["Parameters", "Result-1", "Result-2"],
                interactive=False,
                column_widths=["30%", "30%", "30%"],
                wrap=True
            )

        load_btn_1.click(
            fn=render_result_1,
            inputs=[result_path_1, compared_results],
            outputs=compared_results,
        )
        load_btn_2.click(
            fn=render_result_2,
            inputs=[result_path_2, compared_results],
            outputs=compared_results,
        )

    demo.launch()