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, 'autogbt.sampler.MajorityUnderSampler', 'MajorityUnderSampler', ([], {}), False, 'from autogbt.sampler import MajorityUnderSampler\n'), (30, 'autogbt.sampler.MajorityUnderSampler', 'MajorityUnderSampler', ([], {}), False, 'from autogbt.sampler import MajorityUnderSampler\n'), (20, 'numpy.ones', 'np.ones', (['(10000)'], {}), True, 'import numpy as np\n'), (20, 'numpy.zeros', 'np.zeros', (['(100000)'], {}), True, 'import numpy as np\n'), (14, 'numpy.ones', 'np.ones', (['(10000)'], {}), True, 'import numpy as np\n'), (14, 'numpy.zeros', 'np.zeros', (['(100000)'], {}), True, 'import numpy as np\n'), (26, 'numpy.ones', 'np.ones', (['(10000)'], {}), True, 'import numpy as np\n'), (27, 'numpy.ones', 'np.ones', (['(10000)'], {}), True, 'import numpy as np\n'), (28, 'numpy.ones', 'np.ones', (['(10000)'], {}), True, 'import numpy as np\n')]