import cupy as cp


n_runs = 10
shape = (512, 256, 256)
axis_cases = [(0, 1, 2),]   # dummy

for dtype in (cp.int64, cp.float32, cp.float64, cp.complex64, cp.complex128):
    if dtype in (cp.float32, cp.float64):
        x = cp.random.random(shape, dtype=dtype)
    elif dtype in (cp.int32, cp.int64):
        x = cp.random.randint(0, 10, size=shape, dtype=dtype)
    else:
        x = cp.random.random(shape).astype(dtype) + 1j * cp.random.random(shape).astype(dtype)
    x_np = cp.asnumpy(x) #move to cpu

    for axis in axis_cases:
        for func in ('argmax', 'argmin'):
                keepdims = False
                print("testing", axis, "+", str(dtype), "+", "keepdims={}".format(keepdims), "+", func, "...")
                start = cp.cuda.Event()
                end = cp.cuda.Event()

                cp.cuda.cub_enabled = False
                w = None
                start.record()
                for i in range(n_runs):
                    w = getattr(x, func)()
                end.record()
                end.synchronize()
                t_cp_disabled = cp.cuda.get_elapsed_time(start, end)

                cp.cuda.cub_enabled = True
                y = None
                start.record()
                for i in range(n_runs):
                    y = getattr(x, func)()
                end.record()
                end.synchronize()
                t_cp_enabled = cp.cuda.get_elapsed_time(start, end)

                z = None
                start.record()
                for i in range(n_runs):
                    z = getattr(x_np, func)()
                end.record()
                end.synchronize()
                t_np = cp.cuda.get_elapsed_time(start, end)

                print("CUB enabled: {}, CUB disabled: {}, numpy: {} (ms for {} runs)\n".format(t_cp_enabled, t_cp_disabled, t_np, n_runs))

                try:
                    assert cp.allclose(w, y)
                except AssertionError:
                    print("**************** RESULTS DO NOT MATCH: CUB & reduction ****************")
                    print(w, y)
                try:
                    assert cp.allclose(y, z)
                except AssertionError:
                    print("**************** RESULTS DO NOT MATCH: CUB & NumPy ****************")
                    print(y, z)
        print()