import torch
import math


def cal_GauProb(mu, sigma, x):
    """
    Return the probability of "data" given MoG parameters "mu" and "sigma".
    
    Arguments:
        mu    (BxGxC) - The means of the Gaussians. 
        sigma (BxGxC) - The standard deviation of the Gaussians.
        x     (BxC)   - A batch of data points (coordinates of position).

    Return:
        probabilities (BxG): The probability of each point in the probability
            of the distribution in the corresponding mu/sigma index.
            (Assume the dimensions of the output are independent to each other.)
    """
    x = x.unsqueeze(1).expand_as(mu) # BxC -> Bx1xC -> BxGxC
    prob = torch.rsqrt(torch.tensor(2*math.pi)) * torch.exp(-((x - mu) / sigma)**2 / 2) / sigma
    return torch.prod(prob, dim=2)