import torch


def census_transform(img: torch.Tensor, kernel_size: int) -> torch.Tensor:
    """Census transform
    
    calculates the census transform of an image of shape [N x C x H x W] with 
    batch size N, number of channels C, height H and width W. If C > 1, the census 
    transform is applied independently on each channel.

    Args:
        img (torch.Tensor): input image as torch.Tensor of shape [H x C x H x W]
        kernel_size (int): [description]

    Raises:
        NotImplementedError: kernel size should be 3 or 5.

    Returns:
        torch.Tensor: census transform of img
    """
 
    assert len(img.size()) == 4
    
    if kernel_size != 3 and kernel_size != 5:
        raise NotImplementedError

    n, c, h, w = img.size()

    # get the center idx of census filter
    census_center = margin = int((kernel_size - 1) / 2)

    # init census container
    census = torch.zeros((n, c, h - kernel_size + 1, w - kernel_size + 1), dtype=(torch.int32), device=(img.device))
    center_points = img[:, :, margin:h - margin, margin:w - margin]
    offsets = [(u, v) for v in range(kernel_size) for u in range(kernel_size) if not u == census_center == v]
    for u, v in offsets:
        census = census * 2 + (img[:, :, v:v + h - kernel_size + 1, u:u + w - kernel_size + 1] >= center_points).int()

    census = torch.nn.functional.pad(census, (margin, margin, margin, margin), mode='constant', value=0)

    return census