from argparse import ArgumentParser
from timeit import default_timer as timer
from time import sleep
from os.path import getsize
from os import system

import h5py
import numpy as np

# Blosc filter number
BLOSC_FILTER = 32001

# Codes for compressors shipped with Blosc
BLOSC_BLOSCLZ = 0
BLOSC_LZ4 = 1
BLOSC_LZ4HC = 2
BLOSC_SNAPPY = 3
BLOSC_ZLIB = 4
BLOSC_ZSTD = 5

# Compressors/options to benchmark
TEST_CASES = [
    ('gzip_4', {
        'compression': 'gzip',
        'compression_opts': 4,
        'shuffle': True}),
    ('gzip_6', {
        'compression': 'gzip',
        'compression_opts': 6,
        'shuffle': True}),
    ('lzf', {
        'compression': 'lzf',
        'shuffle': True}),
    ('blosc_blosclz_6', {
        'compression': BLOSC_FILTER,
        'compression_opts': (0, 0, 0, 0, 6, 1, BLOSC_BLOSCLZ),
        'shuffle': False}),
    ('blosc_blosclz_9', {
        'compression': BLOSC_FILTER,
        'compression_opts': (0, 0, 0, 0, 9, 1, BLOSC_BLOSCLZ),
        'shuffle': False}),
    ('blosc_lz4_6', {
        'compression': BLOSC_FILTER,
        'compression_opts': (0, 0, 0, 0, 6, 1, BLOSC_LZ4),
        'shuffle': False}),
    ('blosc_lz4_9', {
        'compression': BLOSC_FILTER,
        'compression_opts': (0, 0, 0, 0, 9, 1, BLOSC_LZ4),
        'shuffle': False}),
    ('blosc_lz4hc_4', {
        'compression': BLOSC_FILTER,
        'compression_opts': (0, 0, 0, 0, 4, 1, BLOSC_LZ4HC),
        'shuffle': False}),
    ('blosc_lz4hc_6', {
        'compression': BLOSC_FILTER,
        'compression_opts': (0, 0, 0, 0, 6, 1, BLOSC_LZ4HC),
        'shuffle': False})
]

def parse_args():
    parser = ArgumentParser()
    parser.add_argument('--truncate-air', action='store_true', help='set voxels < 0.01 to 0')
    parser.add_argument('--cast', action='store_true', help='cast data to float16 and back')
    parser.add_argument('--scaleoffset', type=int, help='apply the scaleoffset filter with the given parameter')
    parser.add_argument('--dataset', type=str, default='reconstruction', help='name of input dataset')
    parser.add_argument('FILE', help='input HDF5 file')
    return parser.parse_args()

def main():
    args = parse_args()

    print('compressor,ctime_mean(s),ctime_std(s),rtime_mean(s),rtime_std(s),size(B)')

    for name, kwargs in TEST_CASES:
        print(name + ',', end='')

        if args.scaleoffset:
            kwargs['scaleoffset'] = args.scaleoffset

        # Benchmark writing/compression of /tmp/out.hdf5
        times = []
        for i in range(3):
            with h5py.File('/tmp/out.hdf5', 'w') as outfile:
                with h5py.File(args.FILE, 'r') as infile:
                    indata = infile[args.dataset].value

                start = timer()

                if args.cast:
                    indata = indata.astype(np.float16).astype(np.float32)

                if args.truncate_air:
                    indata[indata < 0.01] = 0.0

                outfile.create_dataset(
                    args.dataset,
                    data=indata,
                    **kwargs
                )

                end = timer()

                times.append(end - start)
        print('{:.5f},{:.5f},'.format(np.mean(times), np.std(times)), end='')

        # Benchmark reading/decompression of /tmp/out.hdf5
        times = []
        for i in range(3):
            start = timer()
            with h5py.File('/tmp/out.hdf5', 'r') as infile:
                indata = infile[args.dataset].value
            end = timer()
            times.append(end - start)
        print('{:.5f},{:.5f},'.format(np.mean(times), np.std(times)), end='')
        
        # Last column is output file size in bytes
        print(getsize('/tmp/out.hdf5'))

if __name__ == '__main__':
    main()