import torch


def gaussian_mixture_log_prob(x: torch.Tensor, gmm: dict) -> torch.Tensor:
    """Calculate the log-likelihood of a Gaussian mixture model

    Args:
        x: data to evaluate the weighted log-likelihood (N, d)
        gmm: dictionary of a Gaussian mixture model (K components)

    Return: a N-dimensional log-likelihood vector
    """
    mixture_distribution = torch.distributions.Categorical(probs=gmm['pi'])

    if 'scale_tril' in gmm.keys():
        comp_distribution = torch.distributions.multivariate_normal.MultivariateNormal(
            loc=gmm['mean'],
            scale_tril=gmm['scale_tril']
        )
    elif 'cov' in gmm.keys():
        comp_distribution = torch.distributions.multivariate_normal.MultivariateNormal(
            loc=gmm['mean'],
            covariance_matrix=gmm['cov']
        )
    else:
        raise ValueError('Missing scale_tril or cov')

    gm_distribution = torch.distributions.mixture_same_family.MixtureSameFamily(
        mixture_distribution=mixture_distribution,
        component_distribution=comp_distribution
    )

    return gm_distribution.log_prob(x=x)