# Benchmarks from https://github.com/deshaw/versioned-hdf5/issues/155

import h5py
import numpy as np
from versioned_hdf5 import VersionedHDF5File

dt = np.dtype('double')

def time_resize():
    with h5py.File('foo.h5', 'w') as f:
        vf = VersionedHDF5File(f)
        with vf.stage_version('0') as sv:
            sv.create_dataset('bar', (2, 15220, 2),
                              chunks=(300, 100, 2),
                              dtype=dt, data=np.full((2, 15220, 2), 0, dtype=dt))

    with h5py.File('foo.h5', 'r+') as f:
        vf = VersionedHDF5File(f)
        with vf.stage_version('1') as sv:
            bar = sv['bar']
            bar.resize((3, 15222, 2))

time_resize.timeout = 1200

# Pure HDF5 for comparison
def time_resize_hdf5():
    with h5py.File('foo.h5', 'w') as f:
        f.create_dataset('bar', (2, 15220, 2),
                              chunks=(300, 100, 2),
                              dtype=dt, data=np.full((2, 15220, 2), 0,
                                                     dtype=dt),
                         maxshape=(None, None, None))

    with h5py.File('foo.h5', 'r+') as f:
        bar = f['bar']
        bar.resize((3, 15222, 2))

def time_resize_and_write():
    with h5py.File('foo.h5', 'w') as f:
        vf = VersionedHDF5File(f)
        with vf.stage_version('0') as sv:
            sv.create_dataset('bar', (1, 10, 2),
                              chunks=(600, 2, 4),
                              dtype=dt, data=np.full((1, 10, 2), 0, dtype=dt))

    for i in range(1, 100):
        with h5py.File('foo.h5', 'r+') as f:
            vf = VersionedHDF5File(f)
            with vf.stage_version(str(i)) as sv:
                bar = sv['bar']
                bar.resize((1, (i+1) * 10, 2))
                bar[:, -10:, :] = np.full((1, 10, 2), i, dtype=dt)

time_resize_and_write.timeout = 1200

def time_resize_and_write_hdf5_no_copy():
    with h5py.File('foo.h5', 'w') as f:
        f.create_dataset('bar', (1, 10, 2),
                         chunks=(600, 2, 4),
                         dtype=dt, data=np.full((1, 10, 2), 0, dtype=dt),
                         maxshape=(None, None, None))

    for i in range(1, 100):
        with h5py.File('foo.h5', 'r+') as f:
            bar = f['bar']
            bar.resize((1, (i+1) * 10, 2))
            bar[:, -10:, :] = np.full((1, 10, 2), i, dtype=dt)

def time_resize_and_write_hdf5():
    with h5py.File('foo.h5', 'w') as f:
        f.create_dataset('bar0', (1, 10, 2),
                         chunks=(600, 2, 4),
                         dtype=dt, data=np.full((1, 10, 2), 0, dtype=dt),
                         maxshape=(None, None, None))

    for i in range(1, 100):
        with h5py.File('foo.h5', 'r+') as f:
            bar = f.create_dataset('bar%d' % i, chunks=(600, 2, 4), dtype=dt,
                                   data=f['bar%d' % (i-1)],
                                   maxshape=(None, None, None))
            bar.resize((1, (i+1) * 10, 2))
            bar[:, -10:, :] = np.full((1, 10, 2), i, dtype=dt)
