from caffe2.python import core
from hypothesis import given
import caffe2.python.hypothesis_test_util as hu
import numpy as np


class TestBucketizeOp(hu.HypothesisTestCase):
    @given(
        x=hu.tensor(
            min_dim=1, max_dim=2, dtype=np.float32,
            elements=hu.floats(min_value=-5, max_value=5)),
        **hu.gcs)
    def test_bucketize_op(self, x, gc, dc):
        length = np.random.randint(low=1, high=5)
        boundaries = np.random.randn(length) * 5
        boundaries.sort()

        def ref(x, boundaries):
            bucket_idx = np.digitize(x, boundaries, right=True)
            return [bucket_idx]

        op = core.CreateOperator('Bucketize',
                                 ["X"], ["INDICES"],
                                 boundaries=boundaries)
        self.assertReferenceChecks(gc, op, [x, boundaries], ref)


if __name__ == "__main__":
    import unittest
    unittest.main()