import glob
import os
import sys

import h5py
import uproot

################################

def main(root_files):
    merged_file = uproot.concatenate(root_files)
    if 'Events' in merged_file.fields: merged_file = merged_file['Events']
    output_filepath = os.path.join('/'.join(root_files[0].split('/')[:-1]), "all_merged.h5")
    with h5py.File(output_filepath, "w") as output:
        for dataset_name in merged_file.fields:
            output.create_dataset(dataset_name, data=merged_file[dataset_name])

if __name__ == "__main__":
    assert len(sys.argv) == 2, f"Only 1 argument accepted, dirpath to root files"
    dirpath = sys.argv[1]

    root_files = glob.glob(os.path.join(dirpath, "*.root"))
    main(root_files)