import random
import gradio as gr
from optimum_benchmark.task_utils import (
    TASKS_TO_AUTOMODELS,
    infer_task_from_model_name_or_path,
)
from run import run_benchmark
from configs import (
    get_training_config,
    get_inference_config,
    get_neural_compressor_config,
    get_onnxruntime_config,
    get_openvino_config,
    get_pytorch_config,
)

BACKENDS = ["pytorch", "onnxruntime", "openvino", "neural-compressor"]
BENCHMARKS = ["inference", "training"]


with gr.Blocks() as demo:
    # title text
    gr.HTML("<h1 style='text-align: center'>🤗 Optimum Benchmark UI 🏋️</h1>")
    # explanation text
    gr.Markdown("This is a demo space of [Optimum-Benchmark](https://github.com/huggingface/optimum-benchmark.git).")

    model = gr.Textbox(
        label="model",
        value="bert-base-uncased",
    )
    task = gr.Dropdown(
        label="task",
        value="text-classification",
        choices=list(TASKS_TO_AUTOMODELS.keys()),
    )
    device = gr.Dropdown(
        value="cpu",
        label="device",
        choices=["cpu", "cuda"],
    )
    experiment = gr.Textbox(
        label="experiment_name",
        value=f"experiment_{random.getrandbits(16)}",
    )
    model.submit(fn=infer_task_from_model_name_or_path, inputs=model, outputs=task)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                backend = gr.Dropdown(
                    label="backend",
                    choices=BACKENDS,
                    value=BACKENDS[0],
                )

            with gr.Row() as backend_configs:
                with gr.Accordion(label="Pytorch Config", open=False, visible=True):
                    pytorch_config = get_pytorch_config()
                with gr.Accordion(label="OnnxRunTime Config", open=False, visible=False):
                    onnxruntime_config = get_onnxruntime_config()
                with gr.Accordion(label="OpenVINO Config", open=False, visible=False):
                    openvino_config = get_openvino_config()
                with gr.Accordion(label="Neural Compressor Config", open=False, visible=False):
                    neural_compressor_config = get_neural_compressor_config()

        # hide backend configs based on backend
        backend.change(
            inputs=backend,
            outputs=backend_configs.children,
            fn=lambda value: [gr.update(visible=value == key) for key in BACKENDS],
        )

        with gr.Column():
            with gr.Row():
                benchmark = gr.Dropdown(
                    label="benchmark",
                    choices=BENCHMARKS,
                    value=BENCHMARKS[0],
                )

            with gr.Row() as benchmark_configs:
                with gr.Accordion(label="Inference Config", open=False, visible=True):
                    inference_config = get_inference_config()
                with gr.Accordion(label="Training Config", open=False, visible=False):
                    training_config = get_training_config()

        # hide benchmark configs based on benchmark
        benchmark.change(
            inputs=benchmark,
            outputs=benchmark_configs.children,
            fn=lambda value: [gr.update(visible=value == key) for key in BENCHMARKS],
        )

    button = gr.Button(value="Run Benchmark", variant="primary")
    with gr.Accordion(label="LOGS", open=True, visible=False):
        output = gr.HTML()

    button.click(
        fn=run_benchmark,
        inputs={
            experiment,
            model,
            task,
            device,
            backend,
            benchmark,
            *pytorch_config,
            *openvino_config,
            *onnxruntime_config,
            *neural_compressor_config,
            *inference_config,
            *training_config,
        },
        outputs=output,
        queue=True,
    )
    button.click(
        inputs=[],
        outputs=output.parent,
        fn=lambda: gr.update(visible=True),
    )


demo.queue().launch()