import torch


def make_spatial_kernel(kernel_size, bandwidth, isreshape=True):
    """
    Make 2D square smoothness kernel, i.e. h=w
    k = 1/bandwidth * exp(-(pj-pi)**2/(2*bandwidth**2))
    pj, pi = location of pixel
    """
    assert bandwidth > 0, 'bandwidth of kernel must be > 0'
    assert kernel_size % 2 != 0, 'kernel must be odd'
    p_end = (kernel_size - 1) // 2  #kernel center indices
    X = torch.linspace(-p_end, p_end,
                       steps=kernel_size).expand(kernel_size, kernel_size)
    Y = X.clone().t()
    kernel = torch.exp(-(X**2 + Y**2) / (2 * (bandwidth**2)))
    #! due to the require of paper: j#i, thus when j=i, kernel=0
    kernel[p_end, p_end] = 0
    if isreshape:
        return kernel.view(1, 1, kernel_size, kernel_size)
    return kernel