import numpy as np
import pandas as pd
from autogbt.sampler import MajorityUnderSampler


def _test_sample(y):
    sampler = MajorityUnderSampler()
    idx = sampler.sample(y, 40000, 3.0)
    assert len(idx) == 40000
    assert y[idx].sum() == 10000


def test_sample_with_series():
    y = pd.Series(np.concatenate([np.ones((10000)), np.zeros((100000))]))
    y = y.sample(frac=1.0)
    _test_sample(y)


def test_sample_with_ndarray():
    y = np.concatenate([np.ones((10000)), np.zeros((100000))])
    _test_sample(y)


def test_sample_for_regression():
    y = np.concatenate([
        2*np.ones((10000)),
        1*np.ones((10000)),
        0*np.ones((10000)),
    ])
    sampler = MajorityUnderSampler()
    idx = sampler.sample(y, 0.1, 3.0)
    assert len(idx) == 3000

[((7, 14, 7, 36), 'autogbt.sampler.MajorityUnderSampler', 'MajorityUnderSampler', ({}, {}), '()', False, 'from autogbt.sampler import MajorityUnderSampler\n'), ((30, 14, 30, 36), 'autogbt.sampler.MajorityUnderSampler', 'MajorityUnderSampler', ({}, {}), '()', False, 'from autogbt.sampler import MajorityUnderSampler\n'), ((20, 24, 20, 40), 'numpy.ones', 'np.ones', ({(20, 33, 20, 38): '10000'}, {}), '(10000)', True, 'import numpy as np\n'), ((20, 42, 20, 60), 'numpy.zeros', 'np.zeros', ({(20, 52, 20, 58): '100000'}, {}), '(100000)', True, 'import numpy as np\n'), ((14, 34, 14, 50), 'numpy.ones', 'np.ones', ({(14, 43, 14, 48): '10000'}, {}), '(10000)', True, 'import numpy as np\n'), ((14, 52, 14, 70), 'numpy.zeros', 'np.zeros', ({(14, 62, 14, 68): '100000'}, {}), '(100000)', True, 'import numpy as np\n'), ((26, 10, 26, 26), 'numpy.ones', 'np.ones', ({(26, 19, 26, 24): '10000'}, {}), '(10000)', True, 'import numpy as np\n'), ((27, 10, 27, 26), 'numpy.ones', 'np.ones', ({(27, 19, 27, 24): '10000'}, {}), '(10000)', True, 'import numpy as np\n'), ((28, 10, 28, 26), 'numpy.ones', 'np.ones', ({(28, 19, 28, 24): '10000'}, {}), '(10000)', True, 'import numpy as np\n')]