import torch


def quad_kl_div(pi, gamma, ref):
    """Compute the quadratic entropy (KL^otimes(pi otimes gamma | ref))
    with full plans

    Parameters
    ----------
    pi: first input, torch.Tensor of size [Batch, size_X, size_Y]

    gamma: second input torch.Tensor of size [Batch, size_X, size_Y]

    ref: Reference of the KL entropy to compare with (pi otimes gamma)

    Returns
    -------
    div: torch.Tensor of size [Batch]
    Quadratic Kl divergence between each batch.
    """
    massp, massg = pi.sum(), gamma.sum()
    div = (
            massg * torch.sum(pi * (pi / ref + 1e-10).log())
            + massp
            * torch.sum(gamma * (gamma / ref + 1e-10).log())
            - massp * massg
            + ref.sum() ** 2
    )
    return div