from typing import Tuple
from typing import Union
from typing import List
import torch


def gaussian_heatmap(
    image_size: Tuple[int, int],
    center: Union[Tuple[float, float], List[float]],
    sigma: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """
    Creates a Gaussian blob heatmap for a single keypoint.
    The coordinate system is a left-top corner origin with u going left and v going down.

    Args:
        image_size (Tuple(int,int)): image_size (height, width) ! note convention to match store-order and tensor dimensions
        center (Tuple(int,int)): center coordinate (cX,cY) (U,V) ! note: U,V order
        sigma (torch.Tensor): standard deviation of the gaussian blob

    Returns:
        Torch.Tensor: A tensor with zero background, specified size and a Gaussian heatmap around the center.
    """

    # cast keypoints (center) to ints to make grid align with pixel raster.
    #  Otherwise, the AP metric for  d = 1 will not result in 1
    #  if the gt_heatmaps are used as input.
    u_axis = torch.linspace(0, image_size[1] - 1, image_size[1], device=device) - int(center[0])
    v_axis = torch.linspace(0, image_size[0] - 1, image_size[0], device=device) - int(center[1])
    # create grid values in 2D with x and y coordinate centered aroud the keypoint
    xx, yy = torch.meshgrid(v_axis, u_axis)
    ## create gaussian around the centered 2D grids $ exp ( -0.5 (x**2 + y**2) / sigma**2)$
    heatmap = torch.exp(-0.5 * (torch.square(xx) + torch.square(yy)) / torch.square(sigma))
    return heatmap