import torch


def cube_kernel(cube_dim, kernel_size, dilation=1):
    """
    Returns a kernel sampling grid with angular spacing according to the provided cube dimension (and associated computed angular resolution) of a cube map

    cube_dim: length of side of square face of cube map
    kernel_size: (kh, kw)
    """
    # For convenience
    kh, kw = kernel_size

    cube_res = 1 / cube_dim

    # Build the kernel according to the angular resolution of the cube face
    dx = torch.zeros(kernel_size)
    dy = torch.zeros(kernel_size)
    for i in range(kh):
        cur_i = i - (kh // 2)
        for j in range(kw):
            cur_j = j - (kw // 2)
            dx[i, j] = cur_j * dilation * cube_res
            # Flip sign is because +Y is down
            dy[i, j] = cur_i * dilation * -cube_res

    # Returns the kernel differentials as kh x kw
    return dx, dy