import torch


def neighborhood_kernel(kernel_size: int,
                        groups: int,
                        dtype: torch.dtype = torch.float32,
                        device: torch.device = torch.device("cpu")):
    """
        Returns a neighborhood convolution kernel,
        It is a tensor which can be used to extract the K * K points in the neighborhood of a given pixel
        Typically used as a weight matrix in a conv2d operation

        Parameters
        ----------
        kernel_size : int (K)
            the size of the kernel
        groups : int
            the number of groups to replicate
        dtype : torch.dtype
            The type of the weight tensor
        device : torch.device
            The device on which to create the given kernel

        Returns
        -------
        kernel : torch.Tensor [groups * K * K, K, K]
            A weigh matrix for a convolution product which extracts all points in the neighbor

    """

    neighbors = kernel_size * kernel_size
    weights = torch.zeros(neighbors, 1, kernel_size, kernel_size, dtype=dtype, device=device)
    idx = 1
    for i in range(kernel_size):
        for j in range(kernel_size):
            if i == j and i == kernel_size // 2:
                weights[0, :, i, j] = 1.0
            else:
                weights[idx, :, i, j] = 1.0
                idx += 1

    weights_neighborhood = weights
    weights_neighbors = weights.unsqueeze(0).expand(groups, neighbors, 1, kernel_size, kernel_size).reshape(
        groups * neighbors, 1,
        kernel_size, kernel_size)

    return weights_neighborhood, weights_neighbors