"""Write the HDF5 files the h5rs interop tests run on.

usage: gen_files.py OUTDIR

Writes OUTDIR/{earliest,latest}.h5 (the same content with the oldest and the
newest file-format structures: symbol tables and v1 B-trees vs. v2 object
headers, fractal heaps, v2 B-trees and the chunk indexes of HDF5 1.10+), the
pairs the diff tests compare, and a file with a user block. Prints one JSON
object with the values h5py reads back, for the dump tests.
"""

import json
import os
import sys

import h5py
import numpy as np

out = sys.argv[1]


def content(f, dense):
    f.attrs["title"] = "h5rs test"
    f.attrs["version"] = np.int64(3)
    f.attrs["scale"] = np.array([0.5, 1.5], dtype="f4")
    f["contig"] = np.arange(12, dtype="f8").reshape(3, 4)
    dcpl = h5py.h5p.create(h5py.h5p.DATASET_CREATE)
    dcpl.set_layout(h5py.h5d.COMPACT)
    space = h5py.h5s.create_simple((4,))
    h5py.h5d.create(f.id, b"compact", h5py.h5t.STD_I16LE, space, dcpl).write(
        h5py.h5s.ALL, h5py.h5s.ALL, np.arange(4, dtype="<i2")
    )
    g = f.create_group("grp")
    g.attrs["units"] = "m"
    g.create_dataset(
        "gz", data=np.arange(1000, dtype="i4"), chunks=(100,), compression="gzip",
        shuffle=True, fletcher32=True,
    )
    g.create_dataset("ext1", data=np.arange(50, dtype="u2"), chunks=(8,), maxshape=(None,))
    g.create_dataset(
        "ext2", data=np.arange(60, dtype="f4").reshape(6, 10), chunks=(4, 4), maxshape=(None, None)
    )
    g.create_dataset("fixed", data=np.arange(64, dtype="i8").reshape(8, 8), chunks=(3, 3))
    g.create_dataset("single", data=np.arange(10, dtype="i4"), chunks=(10,))
    g.create_dataset("sparse", shape=(100,), dtype="i4", chunks=(10,), fillvalue=-1)
    g["sparse"][20:30] = 7
    sub = g.create_group("sub")
    sub["scalar"] = np.float32(2.5)
    sub["empty"] = h5py.Empty("f8")
    f["strings"] = np.array([b"ab", b"cde"])
    f["vlstr"] = np.array(["x", "yy", "zzz"], dtype=h5py.string_dtype())
    f["cmp"] = np.array([(1, 2.5), (3, 4.5)], dtype=[("a", "i2"), ("b", ">f4")])
    f.create_dataset(
        "enum", data=np.array([0, 1, 1], dtype="u1"),
        dtype=h5py.enum_dtype({"RED": 0, "GREEN": 1}, basetype="u1"),
    )
    f["be"] = np.arange(5, dtype=">i4")
    f["arr"] = np.array([([1, 2, 3],), ([4, 5, 6],)], dtype=[("v", "3i4")])
    f["named_t"] = np.dtype("i8")
    f["soft"] = h5py.SoftLink("/contig")
    f["dangling"] = h5py.SoftLink("/nowhere")
    f["external"] = h5py.ExternalLink("other.h5", "/x")
    f["hard2"] = g["sub"]
    many = f.create_group("many")
    for i in range(12 if dense else 4):
        many[f"d{i:02}"] = np.int32(i)
        many.attrs[f"a{i:02}"] = i


values = {}
# "latest" under HDF5 2.0 writes datatype messages that libhdf5 1.14 tools
# cannot read, so the newest format is taken as 1.14's.
for libver in ("earliest", "latest"):
    path = os.path.join(out, f"{libver}.h5")
    bounds = ("earliest", "v114") if libver == "earliest" else ("v114", "v114")
    with h5py.File(path, "w", libver=bounds) as f:
        content(f, True)
    with h5py.File(path, "r") as f:
        vals = {}

        def grab(name, obj):
            if isinstance(obj, h5py.Dataset) and obj.dtype.kind in "iuf" and obj.shape is not None:
                vals["/" + name] = obj[()].tolist()

        f.visititems(grab)
        values[libver] = vals

# diff pairs
def small(path, data=None, extra=False, attr=False, shape=(3, 4), dtype="f8"):
    with h5py.File(os.path.join(out, path), "w") as f:
        d = np.arange(12, dtype=dtype).reshape(shape) if data is None else data
        f["d"] = d
        f["g/x"] = np.arange(3)
        if extra:
            f["only_here"] = 1
        if attr:
            f["d"].attrs["u"] = 1


base = np.arange(12, dtype="f8").reshape(3, 4)
small("base.h5")
small("same.h5")
changed = base.copy()
changed[0, 2] += 0.001
changed[2, 3] += 0.001
small("changed.h5", data=changed)
small("extra.h5", extra=True)
small("attr.h5", attr=True)
small("reshaped.h5", shape=(4, 3))
small("int.h5", dtype="i4")

with h5py.File(os.path.join(out, "userblock.h5"), "w", userblock_size=1024) as f:
    f["d"] = np.arange(10)

print(json.dumps(values))