import torch
import math


def elastic_deformation(img: torch.Tensor, sample_mode: str = "bilinear", alpha: int = 50,
                        sigma: int = 12) -> torch.Tensor:
    """
    Performs random elastic deformation to the given Tensor image
    :param img: (torch.Tensor) Input image
    :param sample_mode: (str) Resmapling mode
    :param alpha: (int) Scale factor of the deformation
    :param sigma: (int) Standard deviation of the gaussian kernel to be applied
    """
    # Get image shape
    height, width = img.shape[-2:]
    # Get kernel size
    kernel_size = (sigma * 4) + 1
    # Get mean of gaussian kernel
    mean = (kernel_size - 1) / 2.
    # Make gaussian kernel
    # https://discuss.pytorch.org/t/is-there-anyway-to-do-gaussian-filtering-for-an-image-2d-3d-in-pytorch/12351/7
    x_cord = torch.arange(kernel_size, device=img.device)
    x_grid = x_cord.repeat(kernel_size).view(kernel_size, kernel_size)
    y_grid = x_grid.t()
    xy_grid = torch.stack([x_grid, y_grid], dim=-1)
    gaussian_kernel = (1. / (2. * math.pi * sigma ** 2)) \
                      * torch.exp(-torch.sum((xy_grid - mean) ** 2., dim=-1) / (2. * sigma ** 2))
    gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size, kernel_size)
    gaussian_kernel = gaussian_kernel.repeat(1, 1, 1, 1)
    gaussian_kernel.requires_grad = False
    # Make random deformations in the range of [-1, 1]
    dx = (torch.rand((height, width), dtype=torch.float, device=img.device) * 2. - 1.).view(1, 1, height, width)
    dy = (torch.rand((height, width), dtype=torch.float, device=img.device) * 2. - 1.).view(1, 1, height, width)
    # Apply gaussian filter to deformations
    dx, dy = torch.nn.functional.conv2d(input=torch.cat([dx, dy], dim=0), weight=gaussian_kernel, stride=1,
                                        padding=kernel_size // 2).squeeze(dim=0) * alpha
    # Add deformations to coordinate grid
    grid = torch.stack(torch.meshgrid([torch.arange(height, dtype=torch.float, device=img.device),
                                       torch.arange(width, dtype=torch.float, device=img.device)]),
                       dim=-1).unsqueeze(dim=0).flip(dims=(-1,))
    grid[..., 0] += dx
    grid[..., 1] += dy
    # Convert grid to relative sampling location in the range of [-1, 1]
    grid[..., 0] = 2 * (grid[..., 0] - (height // 2)) / height
    grid[..., 1] = 2 * (grid[..., 1] - (width // 2)) / width
    # Resample image
    img_deformed = torch.nn.functional.grid_sample(input=img[None] if img.ndimension() == 3 else img,
                                                   grid=grid, mode=sample_mode, padding_mode='border',
                                                   align_corners=False)[0]
    return img_deformed