from dask.array.percentile import merge_percentiles
import numpy as np
import cupy as cp
import timeit

def bench(nquantiles, ndatasets, lib):
    scale = 100
    calculated_quantiles = nquantiles//2

    finalq = lib.floor(lib.random.rand(nquantiles) * scale)
    qs = list(lib.floor(lib.random.rand(ndatasets, calculated_quantiles) * scale))
    vals = list(lib.random.rand(ndatasets, calculated_quantiles))
    Ns = lib.ones(ndatasets) *  100
    
    dt = 0.
    repeat = 10
    for i in range(repeat):
        t1 = timeit.default_timer()
        merge_percentiles(finalq, qs, vals, Ns=Ns)
        t2 = timeit.default_timer()
        dt += (t2 - t1)
    return (dt) / repeat

import pandas as pd

results = pd.DataFrame(columns=("nquantiles", "ndatasets", "time"))

for nquantiles in (2, 10, 100):
    for ndatasets in (2, 10, 100):
        results.loc[len(results)] = (nquantiles, ndatasets, bench(nquantiles, ndatasets, cp))

results