import torch


def get_probabilistic_prototypes(Nx: dict) -> dict:
    """Calculate the probabilistic prototypes of iid random variable x_i ~ N_i
    Since c = mean(x), then c ~ N(x; mean(mean_x), mean(cov_x))
    
    Args:
        Nx: dictionary with keys: mean, cov, y

    Return: prototypes in a dictionary with keys: mean, cov
    """
    _, d = Nx['mean'].shape
    cls_idx = torch.unique(input=Nx['y'], return_counts=False)
    C = cls_idx.shape[0]

    prototypes = {
        'mean': torch.empty(size=(C, d), device=Nx['mean'].device),
        'cov': torch.empty(size=(C, d, d), device=Nx['cov'].device)
    }

    for c in range(C):
        idx = Nx['y'] == cls_idx[c] # bool
        num_x = torch.sum(input=idx, dim=0)
        prototypes['mean'][c, :] = torch.sum(input=Nx['mean'][idx], dim=0) / num_x
        prototypes['cov'][c, :, :] = torch.sum(input=Nx['cov'][idx], dim=0) / torch.square(input=num_x)
    
    return prototypes