#######################################################################
# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
# Copyright (c) 2023, PyTables Developers <pytables@pytables.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)
#######################################################################

# Benchmark for comparing speeds of getitem of hyperplanes on a
# multidimensional array and using different backends:
# blosc2, PyTables optimized, PyTables filter, and HDF5
# In brief, each approach has its own strengths and weaknesses.
#
# Usage: pass any argument for testing the in-memory backends.
# Else, only persistent containers will be tested.
#
# Uncomment the desired "chunks" value below and run the script
# to get the results with the given chunk size.
# Blosc2 results are only kept as a reference.
#
# Based on "bench/ndarray/compare_getslice.py" from python-blosc2.

import math
import os
import sys
from time import time

import h5py
import hdf5plugin
import numpy as np
import tables

import blosc2

persistent = True if len(sys.argv) == 1 else False
if persistent:
    print("Testing the persistent backends")
else:
    print("Testing the in-memory backends")

# Dimensions and type properties for the arrays
# 3D
# shape = (1000, 2000, 250)
# chunks = (50, 500, 50)
# blocks = (10, 100, 25)

# 4D
shape = (50, 100, 300, 250)
# (uncomment the desired one)
#chunks = (10, 25, 50, 50) # PARAM: small chunk
chunks = (10, 25, 150, 100) # PARAM: big chunk (fits in 32M L3)
blocks = (3, 5, 10, 20)

# Smaller sizes (for quick testing)
# shape = (100, 200, 250)
# chunks = (50, 50, 50)
# blocks = (10, 10, 25)

# shape = (50, 100, 30, 25)
# chunks = (10, 25, 20,  5)
# blocks = ( 3,  5, 10,  2)

dtype = np.dtype(np.int64)

dset_size = math.prod(shape) * dtype.itemsize
# Compression properties
# (LZ4/8 provides a blocksize which fits in 2M L2,
# see "examples/get_blocksize.c" in C-Blosc2)
clevel = 8
cname = "lz4"
nthreads = 8
filter = blosc2.Filter.SHUFFLE
cparams = {
    "codec": blosc2.Codec.LZ4,
    "clevel": clevel,
    "filters": [filter],
    "filters_meta": [0],
    "nthreads": nthreads,
}
dparams = {"nthreads": nthreads}

blocksize = int(np.prod(blocks)) if blocks else 0

fname_b2nd = None
fname_tables = "tables.h5"
fname_h5py = "h5py.h5"
if persistent:
    fname_b2nd = "compare_getslice.b2nd"
    blosc2.remove_urlpath(fname_b2nd)
    fname_tables = "compare_getslice_tables.h5"
    blosc2.remove_urlpath(fname_tables)
    fname_h5py = "compare_getslice_h5py.h5"
    blosc2.remove_urlpath(fname_h5py)

# Create datasets in different formats
# content = np.random.normal(0, 1, int(np.prod(shape)), dtype=dtype).reshape(shape)
# content = np.linspace(0, 1, int(np.prod(shape)), dtype=dtype).reshape(shape)
# (random numbers reduce the effect of similar values in deeper dimensions)
rng = np.random.default_rng()
size = np.prod(shape)
content = rng.integers(low=0, high=10000, size=size, dtype=dtype).reshape(shape)

print("\nCreating datasets...")
# Create and fill a NDArray
t0 = time()
b2 = blosc2.empty(
    shape, dtype=content.dtype, chunks=chunks, blocks=blocks, urlpath=fname_b2nd, cparams=cparams
)
b2[:] = content
t = time() - t0
speed = dset_size / (t * 2**20)
cratio = b2.schunk.cratio
print(f"Time for filling array (blosc2): {t:.3f} s ({speed:.2f} M/s) ; cratio: {cratio:.1f}x")

# Create and fill an HDF5 array (PyTables)
t0 = time()
filters = tables.Filters(complevel=clevel, complib="blosc2:%s" % cname, shuffle=True)
tables.set_blosc_max_threads(nthreads)
if persistent:
    h5f = tables.open_file(fname_tables, "w")
else:
    h5f = tables.open_file(fname_tables, "w", driver="H5FD_CORE", driver_core_backing_store=0)
h5ca = h5f.create_carray(h5f.root, "carray", filters=filters, chunkshape=chunks, obj=content)
t = time() - t0
speed = dset_size / (t * 2**20)
cratio = dset_size / h5ca.size_on_disk
print(f"Time for filling array (hdf5, tables): {t:.3f} s ({speed:.2f} M/s) ; cratio: {cratio:.1f}x")

# Create and fill an HDF5 array (h5py)
t0 = time()
filters = hdf5plugin.Blosc2(clevel=clevel, cname=cname, filters=hdf5plugin.Blosc2.SHUFFLE)
if persistent:
    h5pyf = h5py.File(fname_h5py, "w")
else:
    h5pyf = h5py.File(fname_h5py, "w", driver="core", backing_store=False)
h5d = h5pyf.create_dataset("dataset", dtype=content.dtype, data=content, chunks=chunks, **filters)
t = time() - t0
speed = dset_size / (t * 2**20)
if persistent:
    num_blocks = os.stat(fname_h5py).st_blocks
    # block_size = os.statvfs(fname_h5py).f_bsize
    size_on_disk = num_blocks * 512
    cratio = dset_size / size_on_disk
    print(f"Time for filling array (hdf5, h5py): {t:.3f} s ({speed:.2f} M/s) ; cratio: {cratio:.1f}x")
else:
    print(f"Time for filling array (hdf5, h5py): {t:.3f} s ({speed:.2f} M/s) ; cratio: Not avail")

# Complete reads
print("\nComplete reads...")
t0 = time()
r = b2[:]
t = time() - t0
speed = dset_size / (t * 2**20)
print(f"Time for complete read (blosc2): {t:.3f} s ({speed:.2f} M/s)")

t0 = time()
r = h5ca[:]
t = time() - t0
speed = dset_size / (t * 2**20)
print(f"Time for complete read (hdf5, tables): {t:.3f} s ({speed:.2f} M/s)")

t0 = time()
r = h5d[:]
t = time() - t0
speed = dset_size / (t * 2**20)
print(f"Time for complete read (hdf5, h5py): {t:.3f} s ({speed:.2f} M/s)")

# Reading by slices
print("\nReading by slices...")
# The coordinates for random planes
planes_idx = np.random.randint(0, min(shape), 100)


def time_slices(dset, idx):
    r = None
    if dset.ndim == 3:
        t0 = time()
        if ndim == 0:
            for i in idx:
                r = dset[i, :, :]
        elif ndim == 1:
            for i in idx:
                r = dset[:, i, :]
        else:
            for i in idx:
                r = dset[:, :, i]
        t = time() - t0
        size = r.size * dset.dtype.itemsize * len(idx)
        return t, size / (t * 2**20)
    elif dset.ndim == 4:
        t0 = time()
        if ndim == 0:
            for i in idx:
                r = dset[i, :, :, :]
        elif ndim == 1:
            for i in idx:
                r = dset[:, i, :, :]
        elif ndim == 2:
            for i in idx:
                r = dset[:, :, i, :]
        else:
            for i in idx:
                r = dset[:, :, :, i]
        t = time() - t0
        size = r.size * dset.dtype.itemsize * len(idx)
        return t, size / (t * 2**20)
    raise ValueError(f"ndim == {dset.ndim} is not supported")


for ndim in range(len(shape)):
    print(f"Slicing in dim {ndim}...")

    # Slicing with blosc2
    t, speed = time_slices(b2, planes_idx)
    print(f"Time for reading with getitem (blosc2): {t:.3f} s ({speed:.2f} M/s)")

    # Slicing with hdf5 (PyTables opt)
    os.environ['BLOSC2_FILTER'] = '0'
    t, speed = time_slices(h5ca, planes_idx)
    print(f"Time for reading with getitem (hdf5, tables opt): {t:.3f} s ({speed:.2f} M/s)")

    # Slicing with hdf5 (PyTables noopt)
    os.environ['BLOSC2_FILTER'] = '1'
    t, speed = time_slices(h5ca, planes_idx)
    print(f"Time for reading with getitem (hdf5, tables noopt): {t:.3f} s ({speed:.2f} M/s)")

    # Slicing with hdf5 (h5py)
    t, speed = time_slices(h5d, planes_idx)
    print(f"Time for reading with getitem (hdf5, h5py): {t:.3f} s ({speed:.2f} M/s)")

h5f.close()
h5pyf.close()
# if persistent:
#     os.remove(fname_b2nd)
#     os.remove(fname_tables)
#     os.remove(fname_h5py)
