import torch


def smooth_gaussian1d(x, kernel_sigma, kernel_width):
    """Smooth a signal along the sequence length axis.

    This function is a replacement for the scipy.ndimage.gaussian1d
    function that works on PyTorch tensors. It applies a Gaussian kernel
    to each position which is equivalent to applying a convolution across
    the sequence with weights equal to that of a Gaussian distribution.
    Each sequence, and each channel within the sequence, is smoothed
    independently.

    Parameters
    ----------
    x: torch.tensor, shape=(n_sequences, n_channels, seq_len)
        A tensor to smooth along the last axis. n_channels must be at
        least 1.

    kernel_sigma: float
        The standard deviation of the Gaussian to be applied.

    kernel_width: int
        The width of the kernel to be applied.

    Returns
    -------
    x_smooth: torch.tensor, shape=(n_sequences, n_channels, seq_len)
        The smoothed tensor.
    """

    meshgrid = torch.arange(kernel_width, dtype=torch.float32,
        device=x.device)

    mean = (kernel_width - 1.) / 2.
    kernel = torch.exp(-0.5 * ((meshgrid - mean) / kernel_sigma) ** 2.0)
    kernel = kernel / torch.sum(kernel)
    kernel = kernel.reshape(1, 1, kernel_width).repeat(x.shape[1], 1, 1)
    return torch.nn.functional.conv1d(x, weight=kernel, groups=x.shape[1], 
        padding='same')