from itertools import product

import cupy as cp
import numpy as np
from cupy import testing
from cupyx.profiler import benchmark

# from cupyx.scipy.signal import symiirorder2 as symiir2_gpu
# from scipy.signal import symiirorder2 as symiir2_cpu

from tqdm import tqdm


input_sizes = [5, 10, 100, 1000, 5000, 10000]
# input_sizes = [5000]
modules = [(cp, cp.unique, 'CuPy'), (np, np.unique, 'NumPy')]
# kwargs = [('dtype', [np.float32, np.float64, np.int32, np.int64])]
kwargs = [('dtype', [np.float32])]
# kwargs = [('extrapolate', [True])]
kwargs = list(product(*[list(product([x[0]], x[1])) for x in kwargs]))
measurement = {}


def gather_time(prof):
    cpu_time = prof.cpu_times.mean() * 1000
    gpu_time = prof.gpu_times.mean() * 1000
    return max(cpu_time, gpu_time)


def input_creation(xp, fn, size, **kwargs):
    size = size[0]
    x = testing.shaped_random((size, size), xp, dtype=kwargs['dtype'], scale=0.1)
    return fn, (x,)


def closure(fn, sig):
    # try:
    return fn(*sig, axis=0)
    # except ValueError:
    #     pass


for kwarg_comb in kwargs:
    kwarg_measurement = measurement.get(kwarg_comb, {})
    measurement[kwarg_comb] = kwarg_measurement
    kwarg_comb = dict(kwarg_comb)
    for size in tqdm(input_sizes):
        size_measurement = kwarg_measurement.get(size, {})
        kwarg_measurement[size] = size_measurement
        for xp, cls, _id in modules:
            b, x = input_creation(xp, cls, (size,), **kwarg_comb)
            prof = benchmark(closure, (b, x), n_repeat=100, n_warmup=1)
            size_measurement[_id] = gather_time(prof)


lines = []
for kwarg_comb in kwargs:
    kwarg_line = ', '.join([f'{x}={y.__name__}' for x, y in kwarg_comb])
    lines.append(f'### `{kwarg_line}`\n')
    lines.append('| Size | CuPy (ms) | NumPy (ms) | Speedup |')
    lines.append('|:----:|:---------:|:----------:|:-------:|')
    kwarg_measurement = measurement[kwarg_comb]
    for size in input_sizes:
        comp = f'{size}'
        size_measurement = kwarg_measurement[size]
        times = []
        for _, _, _id in modules:
            time = size_measurement[_id]
            times.append(time)
            time_info = f'{time:3f}'
            comp = f'{comp} | {time_info}'
        speedup = times[1] / times[0]
        comp = f'{comp} | {speedup:3f}'
        lines.append(f'| {comp} |')
    lines.append('\n')

print('\n'.join(lines))