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

[((125, 147), 'autogbt.sampler.MajorityUnderSampler', 'MajorityUnderSampler', ([], {}), '()\n', (145, 147), False, 'from autogbt.sampler import MajorityUnderSampler\n'), ((687, 709), 'autogbt.sampler.MajorityUnderSampler', 'MajorityUnderSampler', ([], {}), '()\n', (707, 709), False, 'from autogbt.sampler import MajorityUnderSampler\n'), ((462, 476), 'numpy.ones', 'np.ones', (['(10000)'], {}), '(10000)\n', (469, 476), True, 'import numpy as np\n'), ((480, 496), 'numpy.zeros', 'np.zeros', (['(100000)'], {}), '(100000)\n', (488, 496), True, 'import numpy as np\n'), ((317, 331), 'numpy.ones', 'np.ones', (['(10000)'], {}), '(10000)\n', (324, 331), True, 'import numpy as np\n'), ((335, 351), 'numpy.zeros', 'np.zeros', (['(100000)'], {}), '(100000)\n', (343, 351), True, 'import numpy as np\n'), ((592, 606), 'numpy.ones', 'np.ones', (['(10000)'], {}), '(10000)\n', (599, 606), True, 'import numpy as np\n'), ((620, 634), 'numpy.ones', 'np.ones', (['(10000)'], {}), '(10000)\n', (627, 634), True, 'import numpy as np\n'), ((648, 662), 'numpy.ones', 'np.ones', (['(10000)'], {}), '(10000)\n', (655, 662), True, 'import numpy as np\n')]