import torch


def shuffle_tensor(input_tensor, right_shift=1):
    """
    given a batch of tensors
    we shuffle them by right shifting the order of images along the first dimension (N)
    Args:
        input_tensor (4d torch.tensor): a batch of images, dim; NCHW
        right_shift (int, optional): [description]. Defaults to 1. when right_shift == batch_size, then it has no effect.
    """
    # shuffle images in a batch, such that the segmentations do not match anymore.
    batch_size = input_tensor.size(0)
    device = input_tensor.device
    idx = torch.arange(batch_size, device=device)
    right_shift = right_shift % batch_size
    if right_shift > 0:
        idx += right_shift
        idx[-right_shift:] = torch.arange(right_shift, device=device)
        shuffled_image_batch = input_tensor[idx]
        input_tensor = shuffled_image_batch
    else:
        print('no shuffle')
    return input_tensor