import os


def get_all_files(root, file_extension):
    files = []
    for folder, _, fs in os.walk(root):
        for f in fs:
            if f.endswith(file_extension):
                files.append(os.path.join(folder, f))
    return files


def remove_duplicated_targets(frames, verbose=False):
    num_target_removed = 0
    total_num_target = 0

    for entry in frames:
        if entry is None or type(entry) == float:
            continue
        if 'targets' not in entry:
            continue

        new_targets = {}
        for target in entry['targets']:
            unit_id = target['unit_id']
            if unit_id not in new_targets:
                new_targets[unit_id] = target
            else:
                old_tick = int(new_targets[unit_id]['tick'])
                new_tick = int(target['tick'])
                # print('remove: old tick:', old_tick, ' new tick:', new_tick)
                if new_tick > old_tick:
                    new_targets[unit_id] = target

        total_num_target += len(entry['targets'])
        num_target_removed += len(entry['targets']) - len(new_targets)
        entry['targets'] = list(new_targets.values())

    if verbose:
        print('remove duplicated: %s out of %s'
              % (num_target_removed, total_num_target))
    return frames
