import torch


def voxel_grid_coords(grid_shape):
    """
    Find normalized voxel coords corresponding to grid cells
    These computations are based on pytorch3d.ops.cubify
    https://pytorch3d.readthedocs.io/en/stable/_modules/pytorch3d/ops/cubify.html#cubify
    Inputs:
    - grid_shape: 3D integer array like containing (D, W, H)
    Returns:
    - FloatTensor of shape (D, W, H, 3)
        xyz (HWD) normalized coordinates for each cell
    """
    assert(len(grid_shape) == 3)
    zyx_grid = torch.meshgrid([torch.arange(grid_shape[i]) for i in range(3)])
    zyx_grid = torch.stack(zyx_grid, dim=0)
    grid_shape_expanded = torch.tensor(grid_shape) \
                               .view(3, 1, 1, 1).expand(-1, *grid_shape) \
                               .float()
    # map to [-1, 1]^3 range from [0, N-1)^3
    zyx_grid = zyx_grid * 2.0 / (grid_shape_expanded - 1.0) - 1.0
    zz, yy, xx = torch.unbind(zyx_grid, dim=0)

    points = torch.stack([xx, yy, zz], dim=-1)
    return points