"""Generate the Blosc2 frames the `filters_blosc2` unit tests decode.

Each case is `<name>.b2f` (a Blosc2 contiguous frame, what the HDF5 Blosc2
filter stores per chunk) and `<name>.out` (what decoding it must give: the
first chunk of a plain frame, or the whole array in C order for a B2ND
frame), or `<name>.err` (a frame clawhdf5 must refuse; the file holds a word
the error message must contain).

These cover what files written by h5py + hdf5plugin never contain, but a
Blosc2 frame may: special chunks (repeated value, NaN, uninitialised), the
delta filter over many blocks and odd type sizes, bit shuffle of blocks that
are not a multiple of 8 elements, shuffle with a byte-group size, forced
stream splitting, multi-chunk B2ND arrays with padded edge chunks and a
chunk of zeros, and features clawhdf5 refuses (dictionaries, registered
filters).

Written with python-blosc2 4.13.1 (c-blosc2 3.3.4) in a scratch venv
(`pip install blosc2`). Re-run only to regenerate:

    python generate.py <this directory>
"""
import os
import sys

import blosc2
import numpy as np

out = sys.argv[1]


def save(name, frame, expected):
    with open(os.path.join(out, name + ".b2f"), "wb") as f:
        f.write(frame)
    with open(os.path.join(out, name + ".out"), "wb") as f:
        f.write(expected)


def save_err(name, frame, word):
    with open(os.path.join(out, name + ".b2f"), "wb") as f:
        f.write(frame)
    with open(os.path.join(out, name + ".err"), "w") as f:
        f.write(word)


def plain(data, **cparams):
    """A one-chunk super-chunk frame of `data`, as hdf5-blosc2 writes."""
    data = np.ascontiguousarray(data)
    cp = blosc2.CParams(typesize=data.dtype.itemsize, **cparams)
    sc = blosc2.SChunk(chunksize=data.nbytes, cparams=cp)
    sc.append_data(data)
    return sc.to_cframe(), data.tobytes()


def special(nitems, dtype, kind, value=None):
    dt = np.dtype(dtype)
    sc = blosc2.SChunk(chunksize=nitems * dt.itemsize,
                       cparams=blosc2.CParams(typesize=dt.itemsize))
    sc.fill_special(nitems, kind, value)
    return sc.to_cframe()


# Special chunks. A repeated value stays in the frame as a 33+ byte chunk;
# NaN and uninitialised chunks become special offsets.
save("value_i4", special(300, "<i4", blosc2.SpecialValue.VALUE, 123456),
     np.full(300, 123456, "<i4").tobytes())
save("value_f8", special(250, "<f8", blosc2.SpecialValue.VALUE, -2.5),
     np.full(250, -2.5, "<f8").tobytes())
save("nan_f4", special(500, "<f4", blosc2.SpecialValue.NAN),
     np.full(500, np.nan, "<f4").tobytes())
save("nan_f8", special(300, "<f8", blosc2.SpecialValue.NAN),
     np.full(300, np.nan, "<f8").tobytes())
save("zero_u2", special(2000, "<u2", blosc2.SpecialValue.ZERO), bytes(4000))
# Uninitialised values: libhdf5 would hand back whatever memory it had;
# clawhdf5 returns zeros.
save("uninit_i8", special(64, "<i8", blosc2.SpecialValue.UNINIT), bytes(512))

rng = np.random.default_rng(11)
ramp = lambda n, dt: ((np.arange(n) * 7) % 1000 + rng.integers(0, 3, n)).astype(dt)
# Slowly varying: what the delta filter is for (noise would be stored raw).
smooth = lambda n, dt: (np.arange(n) // 3 + 1000).astype(dt)

# Delta over many blocks, for type sizes 1, 2, 4, 8, 3 (bytes) and 16 (u64
# pairs).
for dt, n, codec in [("<u1", 2000, blosc2.Codec.LZ4), ("<i2", 1000, blosc2.Codec.BLOSCLZ),
                     ("<i4", 700, blosc2.Codec.LZ4), ("<u8", 400, blosc2.Codec.BLOSCLZ)]:
    save(f"delta_{np.dtype(dt).name}_{codec.name.lower()}",
         *plain(smooth(n, dt), codec=codec, blocksize=256,
                filters=[blosc2.Filter.DELTA], filters_meta=[0]))
rec3 = np.frombuffer(smooth(3 * 300, "<u1").tobytes(), dtype="V3")
save("delta_v3", *plain(rec3, blocksize=300, filters=[blosc2.Filter.DELTA], filters_meta=[0]))
rec16 = np.frombuffer(smooth(2 * 200, "<u8").tobytes(), dtype="V16")
save("delta_shuffle_v16", *plain(rec16, blocksize=512,
                                 filters=[blosc2.Filter.DELTA, blosc2.Filter.SHUFFLE],
                                 filters_meta=[0, 0]))

# Bit shuffle of blocks whose element count is not a multiple of 8 (44-byte
# blocks of 4-byte elements: 8 transposed, 3 copied).
save("bitshuffle_odd_blocks", *plain(ramp(500, "<i4"), codec=blosc2.Codec.ZSTD, blocksize=44,
                                     filters=[blosc2.Filter.BITSHUFFLE], filters_meta=[0]))
# Shuffle in groups of 2 bytes of an 8-byte type (filters_meta).
save("shuffle_meta2", *plain(ramp(500, "<i8"), codec=blosc2.Codec.ZLIB,
                             filters=[blosc2.Filter.SHUFFLE], filters_meta=[2]))
# Streams split per byte, and never split.
save("always_split", *plain(ramp(1000, "<f4"), codec=blosc2.Codec.LZ4HC,
                            splitmode=blosc2.SplitMode.ALWAYS_SPLIT))
save("never_split", *plain(ramp(1500, "<u2"), codec=blosc2.Codec.ZSTD,
                           splitmode=blosc2.SplitMode.NEVER_SPLIT))

# B2ND arrays of several chunks whose edge chunks and blocks are padded, and
# one whose middle chunk is all zeros (a special offset).
for name, shape, chunks, blocks, dt in [
    ("b2nd_2d", (37, 29), (10, 16), (4, 6), "<i4"),
    ("b2nd_3d", (9, 10, 7), (4, 5, 3), (3, 2, 2), "<f4"),
    ("b2nd_4d", (5, 7, 5, 6), (3, 2, 5, 4), (2, 2, 3, 3), "<u2"),
]:
    a = ramp(int(np.prod(shape)), dt).reshape(shape)
    arr = blosc2.asarray(a, chunks=chunks, blocks=blocks)
    save(name, arr.to_cframe(), a.tobytes())
a = ramp(60 * 20, "<i2").reshape(60, 20)
a[20:40, :] = 0
arr = blosc2.asarray(a, chunks=(20, 20), blocks=(8, 16))
save("b2nd_zero_chunk", arr.to_cframe(), a.tobytes())

# Refused: a dictionary, and a registered filter (bytedelta).
frame, _ = plain(ramp(4000, "<i4"), codec=blosc2.Codec.ZSTD, use_dict=True, blocksize=2048)
save_err("zstd_dict", frame, "dictionar")
frame, _ = plain(ramp(500, "<i4"), filters=[blosc2.Filter.SHUFFLE, blosc2.Filter.BYTEDELTA],
                 filters_meta=[0, 4])
save_err("bytedelta", frame, "filter 35")