"""Merge the npz files generated by distributed setting."""
import os
import argparse
import numpy as np


def merge_distributed_npz(data_root):
    """Merge npz files generated by different distributed process.
    Should goes to the npz directory and
        name is '...success_'.
    """
    sub_path = data_root.split('/')
    name = sub_path[-1]
    folder = data_root[:data_root.rindex(name)]
    assert name.endswith('success_')
    all_files = os.listdir(folder)
    # find every file that starts with name
    # accumulate their data and success_rate
    files = [
        file for file in all_files if
        (file.startswith(name) and 'rank' in file)
    ]
    files.sort()
    all_pc, all_lbl, all_target = [], [], []
    for file in files:
        one_file = os.path.join(folder, file)
        npz = np.load(one_file)
        all_pc.append(npz['test_pc'])
        all_lbl.append(npz['test_label'])
        all_target.append(npz['target_label'])
    # concat data
    all_pc = np.concatenate(all_pc, axis=0)
    all_lbl = np.concatenate(all_lbl, axis=0)
    all_target = np.concatenate(all_target, axis=0)
    # sum up attack success rate
    success_rate = 0.
    for file in files:
        idx = file.rindex('-rank_')
        rate = float(file[len(name):idx])
        success_rate += rate
    save_name = name + '{:.4f}.npz'.format(success_rate)
    np.savez(os.path.join(folder, save_name),
             test_pc=all_pc.astype(np.float32),
             test_label=all_lbl.astype(np.uint8),
             target_label=all_target.astype(np.uint8))
    print('data shape', all_pc.shape, all_lbl.shape, all_target.shape)
    print('save to {}'.format(save_name))
    # remove _rank files
    for file in files:
        one_file = os.path.join(folder, file)
        os.remove(one_file)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Point Cloud Recognition')
    parser.add_argument('--data_root', type=str, required=True,
                        default='attack/results')
    args = parser.parse_args()
    merge_distributed_npz(args.data_root)
