import glob
import os
import sys

import h5py
import numpy as np
import uproot

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

def main(root_files, min_valid_targets):
    merged_file = uproot.concatenate(root_files)
    if 'Events' in merged_file.fields: merged_file = merged_file['Events']

    n_tops = max([int(target.split('/')[1].split('t')[-1]) for target in merged_file.fields if 'TARGETS' in target])
    recos = ['FR', 'FB', 'SRqq']
    n_valid = np.zeros(np.array(merged_file[merged_file.fields[0]]).shape[0], dtype=int)
    for i in range(1, n_tops+1):
        ith_any_valid = np.zeros_like(n_valid, dtype=bool)
        for reco in recos:
            ith_any_valid = np.logical_or(ith_any_valid, merged_file[f'TARGETS/{reco}t{i}/MASK'])
        n_valid += np.array(ith_any_valid).astype(int)
    event_mask = (n_valid >= min_valid_targets)

    output_filepath = os.path.join('/'.join(root_files[0].split('/')[:-1]), "all_merged.h5")
    with h5py.File(output_filepath, "w") as output:
        for i, dataset_name in enumerate(merged_file.fields):
            if i == 0: print(f"N events = {np.array(merged_file[dataset_name][event_mask]).shape[0]} ({np.array(merged_file[dataset_name]).shape[0]} events before target masking)")
            output.create_dataset(dataset_name, data=merged_file[dataset_name][event_mask])

if __name__ == "__main__":
    assert len(sys.argv) == 3, f"Only 2 arguments accepted, dirpath to root files and min number of valid targets per top"
    dirpath = sys.argv[1]
    min_valid_targets = int(sys.argv[2])

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