import torch
import math


def solve_for_state(x, logw, loss, n):
    
    """
    Use Geomloss to find the transportation potentials for (x, w) to (w, 1/N)

    :param x: torch.Tensor[N, D]
        The input
    :param logw: torch.Tensor[N]
        The degenerate logweights
    :param loss: geomloss.SamplesLoss
        Needs to be biased for the potentials to correspond to a proper plan
    :param n: int
    :return: torch.Tensor[N], torch.Tensor[N]
        the potentials
    """
    uniform_weights = torch.full_like(logw, math.log(1 / n))
    alpha, beta = loss(uniform_weights, x, logw, x)
    return alpha, beta