import torch


def index2points(points, idx):
    """Construct edge feature for each point

    Parameters
    ----------
    points: (batch_size, num_dims, num_points)
    idx: (batch_size, num_points) or (batch_size, num_points, k)

    Returns
    -------
    edge_features: (batch_size, num_dims, num_points) or (batch_size, num_dims, num_points, k)
    """

    B, C, N = points.shape
    idx_shape = idx.shape

    idx_base = torch.arange(0, B, device=points.device).view(-1, *[1]*(len(idx_shape)-1)) * N # if len(idx_shape) = 3, .view(-1, 1, 1)
    idx = idx + idx_base
    idx = idx.view(-1)

    # (batch_size, num_dims, num_points) -> (batch_size, num_points, num_dims)
    points = points.transpose(2, 1).contiguous()

    feature = points.view(B*N, -1)
    feature = feature[idx, :]
    edge_feature = feature.view(*idx_shape, C)

    # (batch_size, num_points, num_dims) -> (batch_size, num_dims, num_points, ...)
    edge_feature = edge_feature.permute(0, -1, *range(1, len(idx_shape)))
    edge_feature = edge_feature.contiguous()

    return edge_feature