import gradio
import gradio as gr
import plotly.graph_objects as go
from datasets import load_dataset
from huggingface_hub import list_datasets

pipelines = [d.id[20:-21] for d in list_datasets(author='open-source-metrics') if 'checkpoint-downloads' in d.id]


def plot(library: str, stacked: bool):
    dataset = load_dataset(f"open-source-metrics/{library}-checkpoint-downloads")['train']

    dates = dataset['dates']
    axis = dataset.column_names
    axis.remove('dates')

    fig = go.Figure()
    for i in axis:
        fig.add_trace(
            go.Scatter(x=dates, y=dataset[i], mode='lines+markers', name=i, stackgroup='one' if stacked else None)
        )

    fig.show()
    return fig


with gr.Blocks() as demo:
    inputs = [gr.Dropdown(pipelines), gr.Checkbox(label='Stacked')]
    with gr.Row():
        outputs = [gr.Plot()]

    submit = gr.Button('Submit')
    submit.click(fn=plot, inputs=inputs, outputs=outputs)

demo.launch()