import torch


def compute_distance_histograms(a, dx, b, dy):
    """
    Computes the squared distance between histograms of distance of two
    mm-spaces. The histograms are obtained by normalising measures to be
    probabilities.

    Parameters
    ----------
    a: torch.Tensor of size [Batch, size_X]
    Input measure of the first mm-space.

    dx: torch.Tensor of size [Batch, size_X, size_X]
    Input metric of the first mm-space.

    b: torch.Tensor of size [Batch, size_Y]
    Input measure of the second mm-space.

    dy: torch.Tensor of size [Batch, size_X, size_X]
    Input metric of the first mm-space.

    Returns
    -------
    lcost: torch.Tensor of size [size_X, size_Y]
    distances between metric histograms
    """
    h_x = torch.einsum('ij, j->i', dx, a / a.sum())
    h_y = torch.einsum('ij, j->i', dy, b / b.sum())
    lcost = (h_x ** 2)[:, None] + (h_y ** 2)[None, :]
    lcost = lcost - 2 * h_x[:, None] * h_y[None, :]
    return lcost