import pytest
import numpy as np
from rlberry.utils.metrics import metric_lp


@pytest.mark.parametrize("dim", [1, 2, 3, 4])
def test_metrics(dim):
    y = np.zeros(dim)
    x = np.ones(dim)
    scaling_1 = np.ones(dim)
    scaling_2 = 0.5 * np.ones(dim)

    for p in range(1, 10):
        assert np.abs(metric_lp(x, y, p, scaling_1) - np.power(dim, 1.0 / p)) < 1e-15
        assert (
            np.abs(metric_lp(x, y, p, scaling_2) - 2 * np.power(dim, 1.0 / p)) < 1e-15
        )
