import numpy as np
import pytest

from sarus_statistics.ops.mean.local import mean

np.random.seed(0)
NOISE = 1e-9


def test_mean(ops_data, admin_cols):
    public, user_col, weights = admin_cols
    mean_value = mean(
        ops_data,
        data_col='integer_pv',
        user_col=user_col,
        private_col=public,
        weight_col=weights,
        noise=NOISE,
        bounds=(0, 10),
        max_multiplicity=10,
    )

    assert pytest.approx(mean_value, 1e-2) == ops_data['integer_pv'].mean()
    mean_value = mean(
        ops_data,
        data_col='integer_pv',
        user_col=user_col,
        private_col=public,
        weight_col=weights,
        noise=NOISE,
        bounds=(0, 10),
        max_multiplicity=1,
    )

    assert (
        pytest.approx(mean_value, 1e-2)
        == ops_data['integer_pv'].head(100).mean() / 10
    )