import sys
from pathlib import Path

import evaluate
import gradio as gr
import polars as pl
from evaluate import parse_readme

metric = evaluate.load("Aye10032/top5_error_rate")


def compute(data):
    result = {
        "predictions": [list(map(float, pred.split(","))) for pred in data["predictions"]],
        "references": data["references"].cast(pl.Int64).to_list()
    }
    return metric.compute(**result)


local_path = Path(sys.path[0])

default_value = pl.DataFrame({
    'predictions': ['0.82,0.95,0.6,0.14,0.15,0.70', '0.67,0.31,0.01,0.60,0.44,0.51', '0.57,0.06,0.69,0.07,0.96,0.72'],
    'references': ['1', '3', '1']
})

iface = gr.Interface(
    fn=compute,
    inputs=gr.Dataframe(
        headers=['predictions', 'references'],
        col_count=2,
        row_count=1,
        datatype='str',
        type='polars',
        value=default_value
    ),
    outputs=gr.Textbox(label=metric.name),
    description=(
            metric.info.description
            + "\nIf this is a text-based metric, make sure to wrap you input in double quotes."
              " Alternatively you can use a JSON-formatted list as input."
    ),
    title=f"Metric: {metric.name}",
    article=parse_readme(local_path / "README.md"),
)

iface.launch()