def wasserstein_distance(
    u_values: np.ndarray,
    v_values: np.ndarray,
    u_weights: np.ndarray | None = None,
    v_weights: np.ndarray | None = None,
) -> np.ndarray:
    """
    Compute the first Wasserstein distance between two 1D distributions.

    :param u_values: Values of first distribution with shape (nb_samples, feature_dim_1, ..., feature_dim_n).
    :param v_values: Values of second distribution with shape (nb_samples, feature_dim_1, ..., feature_dim_n).
    :param u_weights: Weight for each value. If None, equal weights will be used.
    :param v_weights: Weight for each value. If None, equal weights will be used.
    :return: The Wasserstein distance between the two distributions.
    """
    import scipy

    assert u_values.shape == v_values.shape
    if u_weights is not None:
        assert v_weights is not None
    if u_weights is None:
        assert v_weights is None
    if u_weights is not None and v_weights is not None:
        assert u_weights.shape == v_weights.shape
    if u_weights is not None:
        assert u_values.shape[0] == u_weights.shape[0]

    u_values = u_values.flatten().reshape(u_values.shape[0], -1)
    v_values = v_values.flatten().reshape(v_values.shape[0], -1)

    if u_weights is not None and v_weights is not None:
        u_weights = u_weights.flatten().reshape(u_weights.shape[0], -1)
        v_weights = v_weights.flatten().reshape(v_weights.shape[0], -1)

    w_d = np.zeros(u_values.shape[0])

    for i in range(u_values.shape[0]):
        if u_weights is None and v_weights is None:
            w_d[i] = scipy.stats.wasserstein_distance(u_values[i], v_values[i])
        elif u_weights is not None and v_weights is not None:
            w_d[i] = scipy.stats.wasserstein_distance(u_values[i], v_values[i], u_weights[i], v_weights[i])

    return w_d

    Compute the first Wasserstein distance between two 1D distributions.

    :param u_values: Values of first distribution with shape (nb_samples, feature_dim_1, ..., feature_dim_n).
    :param v_values: Values of second distribution with shape (nb_samples, feature_dim_1, ..., feature_dim_n).
    :param u_weights: Weight for each value. If None, equal weights will be used.
    :param v_weights: Weight for each value. If None, equal weights will be used.
    :return: The Wasserstein distance between the two distributions.
    
art/metrics/metrics.py