import numpy as np
import pytest

from sarus_statistics.ops.exp_quantile.local import (
    get_quantile as get_quantile_exp,
)
from sarus_statistics.ops.joint_exp.local import get_quantiles

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


def test_jointexp(ops_data, admin_cols):
    public, user_col, weights = admin_cols
    quantiles = get_quantiles(
        ops_data,
        data_col='integer',
        user_col=user_col,
        private_col=public,
        weight_col=weights,
        noise=NOISE,
        sampling_ratio=1,
        nb_quantiles=NB_QUANTILES,
        swap=False,
        bounds=(0, 1000),
        max_multiplicity=10,
    )

    assert pytest.approx(list(quantiles.values()), 1e-2) == np.quantile(
        ops_data['integer'], [i / 10 for i in range(1, 10)]
    )


def test_exp_quantile(ops_data, admin_cols):
    public, user_col, weights = admin_cols

    quantile = get_quantile_exp(
        ops_data,
        data_col='integer',
        user_col=user_col,
        private_col=public,
        weight_col=weights,
        noise=NOISE,
        sampling_ratio=1,
        quantile=0.5,
        swap=False,
        bounds=(0, 1000),
        max_multiplicity=10,
    )

    assert pytest.approx(quantile, 1e-2) == np.quantile(
        ops_data['integer'], 0.5
    )