#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Code that saves the histograms (generated by hits_to_histograms) and the mc event_info to a h5 file"""

import h5py

def store_histograms_as_hdf5(hists, mc_infos, filepath_output, compression=(None, None)):
    """
    Takes numpy histograms ('images') for a certain projection as well as the mc_info ('tracks') and saves them to a h5 file.

    Parameters
    ----------
    hists : ndarray(ndim=2)
        Array that contains all histograms for a certain projection.
    mc_infos : ndarray(ndim=2)
        2D array containing important MC information for each event_id, [event_id, particle_type, energy, isCC, categorical event types].
    filepath_output : str
        Complete filepath of the to be created h5 file.
    compression : tuple(None/str, None/int)
        Tuple that specifies if a compression filter should be used. Either ('gzip', 1-9) or ('lzf', None).

    """
    f = h5py.File(filepath_output, 'w')

    chunks_hists = (32,) + hists.shape[1:] if compression[0] is not None else None
    chunks_mc_infos = (32,) + mc_infos.shape[1:] if compression[0] is not None else None
    fletcher32 = True if compression[0] is not None else False

    f.create_dataset('x', data=hists, dtype='uint8', fletcher32=fletcher32, chunks=chunks_hists,
                     compression=compression[0], compression_opts=compression[1], shuffle=False)
    f.create_dataset('y', data=mc_infos, dtype='float32', fletcher32=fletcher32, chunks=chunks_mc_infos,
                     compression=compression[0], compression_opts=compression[1], shuffle=False)

    f.close()