#######################################################################
# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
# All rights reserved.
#
# This source code is licensed under a BSD-style license (found in the
# LICENSE file in the root directory of this source tree)
#######################################################################

# Compute reductions for different array sizes, using the jit decorator
# and different operands (NumPy and NDArray).  Different compression
# levels and codecs can be selected.  For in-memory testing, set ooc to False.

from time import time
import blosc2
import numpy as np
import sys

niter = 5
#dtype = np.dtype("float32")
dtype = np.dtype("float64")
clevel = 1
numpy = False
numpy_jit = False
ooc = True       # for in-memory testing, set to False

cparams = cparams_out = None

# For 64 GB RAM
# sizes_numpy = (1, 5, 10, 20, 30, 35, 40, 45, 50, 55)
# sizes_numpy_jit = (1, 5, 10, 20, 30, 35, 40, 45, 50, 55, 60, 65, 70)
# sizes_clevel0 = (1, 5, 10, 20, 30, 35, 40, 45, 50, 55, 60, 65, 70)
# size_list = (1, 5, 10, 20, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90, 95, 100, 105, 110)  # limit clevel>=1 float64

# For 24 GB RAM
sizes_numpy = (1, 5, 10, 20, 30, 35, 40)  # limit numpy float64
sizes_numpy_jit = (1, 5, 10, 20, 30, 35, 40, 45)  # limit numpy float64
sizes_clevel0 = (1, 5, 10, 20, 30, 35, 40, 45)  # limit clevel==0 float64
#sizes_clevel0 = (50, 55, 60, 65, 70)  # extra sizes for clevel==0 float64
size_list = (1, 5, 10, 20, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90)  # limit clevel>=1 float64

codec = "LZ4"  # default codec
if len(sys.argv) > 2:
    codec = sys.argv[2]
if len(sys.argv) > 1:
    try:
        clevel = int(sys.argv[1])
    except ValueError:
        if sys.argv[1] == "numpy":
            numpy = True
        elif sys.argv[1] == "numpy_jit":
            numpy = True
            numpy_jit = True
        else:
            raise ValueError("Invalid argument")

apath = bpath = None
if numpy:
    print("Using NumPy arrays as operands")
else:
    print("Using NDArray arrays as operands")
    cparams = cparams_out = blosc2.CParams(clevel=clevel, codec=blosc2.Codec[codec])
    # cparams_out = blosc2.CParams(clevel=clevel, codec=blosc2.Codec.LZ4)
    print("Using cparams: ", cparams)
    if ooc:
        apath = "a.b2nd"
        bpath = "b.b2nd"


# The reductions to compute
def compute_reduction_numpy(a, b, c):
    return np.sum(((a ** 3 + np.sin(a * 2)) < c) & (b > 0), axis=1)

@blosc2.jit
def compute_reduction(a, b, c):
    return np.sum(((a ** 3 + np.sin(a * 2)) < c) & (b > 0), axis=1)


create_times = []
compute_times = []
for n in size_list:
    if clevel == 0 and n not in sizes_clevel0:
        continue
    if numpy_jit and n not in sizes_numpy_jit:
        continue
    if numpy and not numpy_jit and n not in sizes_numpy:
        continue
    N = n * 1000
    print(f"\nN = {n}000, {dtype=}, size={N ** 2 * 2 * dtype.itemsize / 2**30 }")
    chunks = (100, N)
    blocks = (1, N)
    chunks, blocks = None, None  # automatic chunk and block sizes
    # Lossy compression
    #filters = [blosc2.Filter.TRUNC_PREC, blosc2.Filter.SHUFFLE]
    #filters_meta = [8, 0]  # keep 8 bits of precision in mantissa
    #cparams = blosc2.CParams(clevel=1, codec=blosc2.Codec.LZ4, filters=filters, filters_meta=filters_meta)

    # Create some data operands
    t0 = time()
    if numpy:
        a = np.linspace(0, 1, N * N, dtype=dtype).reshape(N, N)
        b = np.linspace(1, 2, N * N, dtype=dtype).reshape(N, N)
        #b = a + 1
        c = np.linspace(-10, 10, N, dtype=dtype)
    else:
        a = blosc2.linspace(0, 1, N * N, dtype=dtype, shape=(N, N), cparams=cparams, urlpath=apath, mode="w")
        print("a.chunks, a.blocks, a.schunk.cratio: ", a.chunks, a.blocks, a.schunk.cratio)
        b = blosc2.linspace(1, 2, N * N, dtype=dtype, shape=(N, N), cparams=cparams, urlpath=bpath, mode="w")
        #b = (a + 1).compute(cparams=cparams, chunks=chunks, blocks=blocks)
        #print(b.chunks, b.blocks, b.schunk.cratio, b.cparams)
        c = blosc2.linspace(-10, 10, N, dtype=dtype, cparams=cparams)  # broadcasting is supported
        #c = blosc2.linspace(-10, 10, N * N, dtype=dtype, shape=(N, N), cparams=cparams)
    t1 = time() - t0
    print(f"Time to create data: {t1:.4f}")
    create_times.append(t1)

    if numpy:
        if numpy_jit:
            out = compute_reduction(a, b, c)
            t0 = time()
            for i in range(niter):
                out = compute_reduction(a, b, c)
            t1 = (time() - t0) / niter
            print(f"Time to compute with numpy_jit and NumPy operands: {t1:.4f}")
        else:
            t0 = time()
            nout = compute_reduction_numpy(a, b, c)
            t1 = time() - t0
            print(f"Time to compute with NumPy engine: {t1:.4f}")
    else:
        out = compute_reduction(a, b, c)
        t0 = time()
        for i in range(niter):
            out = compute_reduction(a, b, c)
        t1 = (time() - t0) / niter
        print(f"Time to compute with numpy_jit and {clevel=}: {t1:.4f}")
    compute_times.append(t1)
    del a, b, c

print("\nCreate times: [", ", ".join([f"{t:.4f}" for t in create_times]), "]")
print("Compute times: [", ", ".join([f"{t:.4f}" for t in compute_times]), "]")
print("End of run!\n\n")
