import gradio as gr
from gradio_leaderboard import Leaderboard, SelectColumns, ColumnFilter

from src.common.data import load_dataset
from src.common.schema import DatasetSchema
from src.common.paths import DOCS_PATH, DATASET_NAME


_dataset = load_dataset()


def show_dataset_example(id_: str) -> str:
    example = _dataset[_dataset[DatasetSchema.id_] == id_].iloc[0]
    text = f"""
## Task

**{example[DatasetSchema.description]}**

{example[DatasetSchema.task_text]}

### Answer type

```python
{example[DatasetSchema.answer_type]}
```

({example[DatasetSchema.task_note]})

## Answer

{example[DatasetSchema.answer_text]}

### Correct answer

```json
{example[DatasetSchema.correct_answer]}
```

### Answer check

{example[DatasetSchema.check_type]}
""".strip()
    if example[DatasetSchema.check_function]:
        text += f"\n\n#### Check function\n\n```python\n{example[DatasetSchema.check_function]}\n```"
    return text


with gr.Blocks(
    title="ROMB Leaderboard v1.0",
    theme=gr.themes.Ocean(
        primary_hue=gr.themes.colors.green,
    ),
) as application:
    gr.Markdown("# 🥇 ROMB - Russian Olympiad Math Benchmark")
    gr.Markdown(
        f"See ROMB-1.0 dataset there - [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME})."
    )
    with gr.Tabs():
        with gr.Tab("Leaderboard"):
            gr.Markdown("In progress...")
        with gr.Tab("Evaluate"):
            gr.Markdown((DOCS_PATH / "evaluate.md").read_text())
        with gr.Tab("Submit"):
            gr.Markdown("In progress...")
        with gr.Tab("Task example"):
            gr.Interface(
                fn=show_dataset_example,
                inputs=gr.Dropdown(
                    choices=list(_dataset[DatasetSchema.id_]),
                    value=0,
                    label="Example ID",
                ),
                outputs=gr.Markdown(label="Example"),
                title="Show Dataset Example",
                examples=[
                    [8],
                    [14],
                    [17],
                    [22],
                    [40],
                    [230],
                ],
            )


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