from __future__ import annotations

import torch
from torch import Tensor


def get_gaussian_2d(
    kernel_size: int | tuple[int, int], sigma: float | tuple[float, float]
) -> Tensor:
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    if isinstance(sigma, float | int):
        sigma = (float(sigma), float(sigma))

    kx, ky = kernel_size
    sx, sy = sigma

    x = torch.arange(kx).float() - (kx - 1) / 2
    y = torch.arange(ky).float() - (ky - 1) / 2

    gauss_x = torch.exp(-(x**2) / (2 * sx**2))
    gauss_y = torch.exp(-(y**2) / (2 * sy**2))

    gauss_2d = gauss_x.view(1, -1) * gauss_y.view(-1, 1)
    gauss_2d = gauss_2d / gauss_2d.sum()

    return gauss_2d
