import torch


def apply_distance_bound(data_region_i, data_region_i_orig, args):
    """
    Input:
        data_region_i: (S,3) tensor, current region i points
        data_region_i_orig: (S,3) tensor, original region i points
    Return:
        data_region_i: modified data_region_i
        count: number of points that exceed distance bound
    """
    with torch.no_grad():
        region_i_diff = data_region_i - data_region_i_orig #(S,3)
        region_i_diff_distance = torch.norm(region_i_diff, dim=1) #(S,)

        total_points = region_i_diff_distance.shape[0]
        count = 0
        for i in range(total_points):  # check bound
            if region_i_diff_distance[i] > args.dist_threshold:
                count += 1
                data_region_i[i].data = data_region_i_orig[i].data + args.dist_threshold * region_i_diff[i] / region_i_diff_distance[i]
    return data_region_i, count