import torch


def get_shifts(padding, data_dim, pad_type='hypercubic', include_batch=False):
    """
    Arguments:
        padding: number of padding to add to shifts
        data_dim: dimension of data
    Returns
        shifts:
            Tensor of shape ((2 * padding + 1) ** data_dim) x  data_dim
            Each row of shifts represent nearby coordinates

    Ex)
        >>> get_shifts(1, 2)
        torch.tensor([
            [-1, -1], [-1, 0], [-1, 1],
            [0, -1], [0, 0], [0, 1],
            [1, -1], [1, 0], [1, -1]
        ])
    """

    shifts = []
    if pad_type == 'hypercubic':
        for i in range(-padding, padding + 1):
            for j in range(-padding, padding + 1):
                if data_dim == 2:
                    shifts.append([i, j])
                    continue
                for k in range(-padding, padding + 1):
                    shifts.append([i, j, k])
    elif pad_type == 'hypercross':
        for x in range(padding + 1):
            for y in range(padding + 1 - x):
                if data_dim == 2:
                    shifts.append([x, y])
                    shifts.append([-x, y])
                    shifts.append([x, -y])
                    shifts.append([-x, -y])
                    continue
                for z in range(padding + 1 - x - y):
                    shifts.append([x, y, z])
                    shifts.append([-x, y, z])
                    shifts.append([x, -y, z])
                    shifts.append([x, y, -z])
                    shifts.append([-x, -y, z])
                    shifts.append([-x, y, -z])
                    shifts.append([x, -y, -z])
                    shifts.append([-x, -y, -z])
    else:
        raise ValueError('pad_type {} not allowed'.format(pad_type))

    shifts = torch.unique(torch.Tensor(shifts), dim=0)
    if include_batch:
        # shifts = ME.utils.batched_coordinates(shifts)
        shifts = torch.cat([torch.zeros(shifts.shape[0], 1), shifts], dim=1)
    return shifts