import torch
torch.set_printoptions(profile="full")
import cv2

import torch.nn as nn
import torch.nn.functional as F
import numpy as np


def laplace(mask, kernel_width=3):
    # laplace
    kernel_size = 2 * kernel_width + 1
    laplace_kernel = - torch.ones(1, 1, kernel_size, kernel_size)
    laplace_kernel[0, 0, kernel_width, kernel_width] = kernel_size ** 2 - 1
    laplace_kernel = laplace_kernel / (kernel_size ** 2 - 1)
    print(f'kernel: min: {torch.min(laplace_kernel)}, max: {torch.max(laplace_kernel)}')
    # divergence
    divergence = F.conv2d(mask, laplace_kernel, bias=None, stride=1, padding=kernel_width)
    print(f'divergence: min: {torch.min(divergence)}, max: {torch.max(divergence)}')

    img_div = divergence.abs().squeeze().numpy() * 255
    img_div = img_div.astype(np.uint8)
    cv2.imwrite(f'./demo/div_{kernel_width}.jpg', img_div)


def main():
    # mask
    mask = torch.zeros(1, 1, 64, 64)
    mask[:, :, 16:48, 16:48] = 1
    img_mask = mask.squeeze().numpy() * 255
    img_mask = img_mask.astype(np.uint8)
    cv2.imwrite('./demo/mask.jpg', img_mask)
    # laplace
    for kw in [1, 3, 5, 7]:
        laplace(mask, kw)


if __name__ == '__main__':
    main()