import json

import pandas as pd
from datasets import load_dataset


import gradio as gr
from gradio_leaderboard import Leaderboard, SelectColumns, ColumnFilter
import config
from pathlib import Path
import pandas as pd
import random

abs_path = Path(__file__).parent
submissions = json.load(open(abs_path / "submissions.json"))

TASKS = ["mmlu"]


def load_submissions():
    all_submissions = []

    for submission in submissions["submissions"]:
        ds = load_dataset(submission["results-dataset"], "results")
        ds = ds.filter(lambda x: x["task"] in TASKS)

        all_accuracy = []

        for result in ds["train"]:
            submission[result["task"]] = result["accuracy"]
            all_accuracy.append(result["accuracy"])

        submission["Average ⬆️"] = sum(all_accuracy) / len(all_accuracy)

        submission["username"] = (
            f"[{submission['username']}](https://huggingface.co/{submission['username']})"
        )
        submission["model_name"] = (
            f"[{submission['model_name']}](https://huggingface.co/{submission['model_name']})"
        )

        submission["results-dataset"] = (
            f"[🔗](https://huggingface.co/datasets/{submission['results-dataset']})"
        )

        all_submissions.append(submission)

    all_submissions = pd.DataFrame(all_submissions)
    return all_submissions


with gr.Blocks() as demo:
    gr.Markdown("""
    # 🥇 a smol course leaderboad
    
    A leaderboard of smol course students' submissions. 
    """)
    with gr.Tabs():
        with gr.Tab("Demo"):
            Leaderboard(
                value=load_submissions(),
                select_columns=SelectColumns(
                    default_selection=config.ON_LOAD_COLUMNS,
                    cant_deselect=["username", "model_name"],
                    label="Select Columns to Display:",
                ),
                search_columns=["username", "model_name"],
                # hide_columns=["chapter"],
                # filter_columns=[
                #     "username",
                #     "model_name",
                #     # "Precision",
                #     # ColumnFilter("MOE", type="boolean", default=False, label="MoE"),
                #     # ColumnFilter("Flagged", type="boolean", default=False),
                #     # ColumnFilter("#Params (B)", default=[30, 80]),
                # ],
                datatype=config.TYPES,
                # column_widths=["2%", "33%"],
            )
        with gr.Tab("Docs"):
            gr.Markdown((Path(__file__).parent / "docs.md").read_text())

if __name__ == "__main__":
    demo.launch()