from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import numpy as np

from caffe2.python import core, workspace
from hypothesis import given
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.serialized_test.serialized_test_util as serial
import hypothesis.strategies as st
import hypothesis.extra.numpy as hnp

# Basic implementation of gather for axis == 0, shich is lookup of indices
# in the outer dimention. Keeping it for reference here, although is similar
# to more general funciton below.
def ref_gather_axis0():
    def inner(data, ind):
        if ind.size == 0 or data.shape[0] == 0:
            return [np.zeros((0, 10, 20)).astype(np.float32)]
        output = [data[i] for i in ind]
        return [output]
    return inner

# Returns axis-based lookup. We just use numpy take() which handles different
# axis values as we want.
def ref_gather(axis):
    def inner(data, ind):
        if ind.size == 0 or data.shape[axis] == 0:
            shape = list(data.shape)
            shape[0] = 0
            return [np.zeros(tuple(shape)).astype(np.float32)]
        # np.take() does axis lookup same as gather
        output = data.take(ind, axis).astype(np.float32)
        return [output]
    return inner

class TestGatherOps(serial.SerializedTestCase):
    @given(rows_num=st.integers(0, 10000),
           index_num=st.integers(0, 5000),
           **hu.gcs)
    def test_gather_ops(self, rows_num, index_num, gc, dc):
        data = np.random.random((rows_num, 10, 20)).astype(np.float32)

        if rows_num > 0:
            ind = np.random.randint(rows_num, size=(index_num, )).astype('int32')
        else:
            ind = np.random.randint(10, size=(index_num, )).astype('int32')
        op = core.CreateOperator(
            'Gather',
            ['data', 'ind'],
            ['output'])

        self.assertReferenceChecks(gc, op, [data, ind], ref_gather_axis0())
        self.assertDeviceChecks(dc, op, [data, ind], [0])
        return

    # Test axis == 2, this keeps outer dimension but will replace data
    # within axis by lookup of index array (repeated for each outer entry)
    @given(batch_num=st.integers(1, 4000),
        rows_num=st.integers(1, 6),
        index_num=st.integers(1, 20),
        **hu.gcs)
    def test_gather_ops_axis2(self, batch_num, rows_num, index_num, gc, dc):
        data = np.random.random((batch_num, rows_num, 5)).astype(np.float32)
        ind = np.random.randint(5, size=(index_num, )).astype('int32')
        op = core.CreateOperator(
            'Gather',
            ['data', 'ind'],
            ['output'],
            axis=2)

        self.assertReferenceChecks(gc, op, [data, ind], ref_gather(axis=2))
        self.assertDeviceChecks(dc, op, [data, ind], [0])
        return


# Generates data arrays of max dims 10x100x2 and indexing array up to rows_num
@st.composite
def _inputs(draw):
    batch_size = draw(st.integers(2, 10))
    rows_num = draw(st.integers(1, 100))
    block_size = draw(st.integers(1, 2))
    index_num = draw(st.integers(1, 10))
    return (
        draw(hnp.arrays(
            np.float32,
            (batch_size, rows_num, block_size),
            elements=st.floats(-10.0, 10.0),
        )),
        draw(hnp.arrays(
            np.int32,
            (index_num, 1),
            elements=st.integers(0, rows_num - 1),
        )),
    )

class TestBatchGatherOps(hu.HypothesisTestCase):
    @given(inputs=_inputs(),
           **hu.gcs)
    def test_batch_gather_ops(self, inputs, gc, dc):
        data, ind = inputs
        op = core.CreateOperator(
            'BatchGather',
            ['data', 'ind'],
            ['output'])
        self.assertReferenceChecks(gc, op, [data, ind], ref_gather(axis=1))
        self.assertGradientChecks(gc, op, [data, ind], 0, [0])


class TestGatherFused8BitRowwise(hu.HypothesisTestCase):
    @given(rows_num=st.integers(1, 10000),
           cols_num=st.integers(1, 128),
           index_num=st.integers(0, 5000),
           **hu.gcs)
    def test_batch_gather_ops(self, rows_num, cols_num, index_num, gc, dc):
        data = np.random.random((rows_num, cols_num)).astype(np.float32)
        ind = np.random.randint(rows_num, size=(index_num, )).astype('int32')

        net = core.Net("bench")

        quantized_data = net.FloatToFused8BitRowwiseQuantized(
            'data', 'quantized_data')
        dequantized_data = net.Fused8BitRowwiseQuantizedToFloat(
            quantized_data, 'dequantized_data')

        net.Gather(
            [dequantized_data, 'ind'], 'gather_reference')
        net.GatherFused8BitRowwise(
            [quantized_data, 'ind'], 'gather_quantized')

        workspace.FeedBlob('data', data)
        workspace.FeedBlob('ind', ind)
        workspace.CreateNet(net)
        workspace.RunNetOnce(net)

        gather_reference = workspace.FetchBlob('gather_reference')
        gather_quantized = workspace.FetchBlob('gather_quantized')
        np.testing.assert_array_almost_equal(gather_reference, gather_quantized)



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