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


demo = gr.Interface(fn=plot, inputs=[
    gr.Dropdown(pipelines),
    gr.Checkbox(label='Stacked')
], outputs=gr.Plot())

demo.launch()