import torch


def kl_prox_softmin(ecost, a, b, eps, rho, rho2):
    """Prepares functions which perform updates of the Sikhorn algorithm
    in exponential scale.

    Parameters
    ----------
    ecost: torch.Tensor of size [Batch, size_X, size_Y]
    Exponential of the cost. Kernel of Sinkhorn operator.

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

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

    eps: float
    Strength of entropic regularization.

    rho: float
    Strength of penalty on the first marginal of pi.

    rho2: float
    Strength of penalty on the first marginal of pi. If set to None it is
    equal to rho.

    Returns
    ----------
    s_x: callable function
    Map outputing updates of potential from Y to X.

    s_y: callable function
    Map outputing updates of potential from X to Y.
    """

    tau = 1.0 / (1.0 + eps / rho)
    tau2 = 1.0 / (1.0 + eps / rho2)

    def s_y(v):
        return torch.einsum("ij,j->i", ecost, b * v) ** (-tau2)

    def s_x(u):
        return torch.einsum("ij,i->j", ecost, a * u) ** (-tau)

    return s_x, s_y