import os.path

import gradio as gr
import numpy as np
import pandas as pd

from constants import *


def get_download_link_model(task, dataset, example):
    _task_path = TASK_PATH_MAPPING[task]
    _dataset_path = DATASET_PATH_MAPPING[dataset]
    _example_path = EXAMPLE_PATH_MAPPING[example]
    return os.path.join("data", _task_path, _dataset_path, "weight", f"{_example_path}.zip")


def get_download_link_json(task, dataset, example):
    _task_path = TASK_PATH_MAPPING[task]
    _dataset_path = DATASET_PATH_MAPPING[dataset]
    _example_path = EXAMPLE_PATH_MAPPING[example]
    if _task_path == "common":
        return os.path.join("data", _task_path, _dataset_path, "json", f"{_example_path}.jsonl")
    else:
        return os.path.join("data", _task_path, _dataset_path, "json", f"{_example_path}.json")


def get_data(task, dataset, example):
    _task_path = TASK_PATH_MAPPING[task]
    _dataset_path = DATASET_PATH_MAPPING[dataset]
    _example_path = EXAMPLE_PATH_MAPPING[example]
    csv_file = os.path.join("data", _task_path, _dataset_path, "csv", f"{_example_path}.csv")
    if not os.path.exists(csv_file):
        return None
    read_data = pd.read_csv(csv_file)
    data = pd.DataFrame(columns=COLUMN_NAMES)

    if _task_path == "coding":
        for index, row in read_data.iterrows():
            data = data._append({
                "Prompt": row["prompt"],
                "Pass@1": round(float(row["pass@1"]) * 100, 3),
                "Pass@5": round(float(row["pass@5"]) * 100, 3),
                "Pass@10": round(float(row["pass@10"]) * 100, 3),
                "Correctness": "N/A"
            }, ignore_index=True)
    elif _task_path == "common":
        for index, row in read_data.iterrows():
            data = data._append({
                "Prompt": row["prompt"],
                "Pass@1": None,
                "Pass@5": None,
                "Pass@10": None,
                "Correctness": "✅" if row["correctness"] else "❌"
            }, ignore_index=True)
    elif _task_path == "math":
        for index, row in read_data.iterrows():
            data = data._append({
                "Prompt": row["prompt"],
                "Pass@1": None,
                "Pass@5": None,
                "Pass@10": None,
                "Correctness": "✅" if row["correctness"] else "❌"
            }, ignore_index=True)

    return data


with gr.Blocks() as demo_board:
    gr.HTML(DND_HEADER)
    gr.Markdown(DND_INTRODUCTION)

    task = gr.Radio(
        label="Task",
        choices=TASK_LIST,
        value=TASK_LIST[0],
        interactive=True,
    )
    dataset = gr.Radio(
        label="Dataset",
        choices=TASK_DATASET_LIST[task.value],
        value=TASK_DATASET_LIST[task.value][0],
        interactive=True
    )
    example = gr.Radio(
        label="Example",
        choices=EXAMPLE_LIST,
        value=EXAMPLE_LIST[0],
        interactive=True,
    )

    board = gr.components.Dataframe(
        value=get_data(task.value, dataset.value, example.value),
        column_widths=["60%", "10%", "10%", "10%", "10%"],
        headers=COLUMN_NAMES,
        type="pandas",
        datatype=DATA_TITLE_TYPE,
        interactive=False,
        visible=True,
        max_height=500,
    )

    task.change(lambda t: gr.Radio(
        label="Dataset",
        choices=TASK_DATASET_LIST[t],
        value=TASK_DATASET_LIST[t][0],
        interactive=True,
    ), inputs=[task], outputs=dataset)

    for component in [task, dataset, example]:
        component.change(lambda t, d, e: gr.components.Dataframe(
            value=get_data(t, d, e),
            column_widths=["60%", "10%", "10%", "10%", "10%"],
            headers=COLUMN_NAMES,
            type="pandas",
            datatype=DATA_TITLE_TYPE,
            interactive=False,
            visible=True,
            max_height=500,
        ), inputs=[task, dataset, example], outputs=board)

    with gr.Row():
        json_downloader = gr.DownloadButton("Download JSON", visible=True)
        model_downloader = gr.DownloadButton("Download Model", visible=True)
        json_downloader.click(
            fn=get_download_link_json,
            inputs=[task, dataset, example],
            outputs=json_downloader,
        )
        model_downloader.click(
            fn=get_download_link_model,
            inputs=[task, dataset, example],
            outputs=model_downloader,
        )

    citation_button = gr.Textbox(
        value=CITATION_BUTTON_TEXT,
        label=CITATION_BUTTON_LABEL,
        elem_id="citation-button",
        lines=6,
        show_copy_button=True,
    )

demo_board.launch()