from caffe2.python import core
from functools import partial
from hypothesis import given, settings
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.serialized_test.serialized_test_util as serial
import hypothesis.strategies as st
import numpy as np
import unittest
from caffe2.python import workspace


def _gen_test_add_padding(with_pad_data=True,
                          is_remove=False):
    def gen_with_size(args):
        lengths, inner_shape = args
        data_dim = [sum(lengths)] + inner_shape
        lengths = np.array(lengths, dtype=np.int32)
        if with_pad_data:
            return st.tuples(
                st.just(lengths),
                hu.arrays(data_dim),
                hu.arrays(inner_shape),
                hu.arrays(inner_shape))
        else:
            return st.tuples(st.just(lengths), hu.arrays(data_dim))

    min_len = 4 if is_remove else 0
    lengths = st.lists(
        st.integers(min_value=min_len, max_value=10),
        min_size=0,
        max_size=5)
    inner_shape = st.lists(
        st.integers(min_value=1, max_value=3),
        min_size=0,
        max_size=2)
    return st.tuples(lengths, inner_shape).flatmap(gen_with_size)


def _add_padding_ref(
        start_pad_width, end_pad_width, ret_lengths,
        data, lengths, start_padding=None, end_padding=None):
    if start_padding is None:
        start_padding = np.zeros(data.shape[1:], dtype=data.dtype)
    end_padding = (
        end_padding if end_padding is not None else start_padding)
    out_size = data.shape[0] + (
        start_pad_width + end_pad_width) * len(lengths)
    out = np.ndarray((out_size,) + data.shape[1:])
    in_ptr = 0
    out_ptr = 0
    for length in lengths:
        out[out_ptr:(out_ptr + start_pad_width)] = start_padding
        out_ptr += start_pad_width
        out[out_ptr:(out_ptr + length)] = data[in_ptr:(in_ptr + length)]
        in_ptr += length
        out_ptr += length
        out[out_ptr:(out_ptr + end_pad_width)] = end_padding
        out_ptr += end_pad_width
    lengths_out = lengths + (start_pad_width + end_pad_width)
    if ret_lengths:
        return (out, lengths_out)
    else:
        return (out, )


def _remove_padding_ref(start_pad_width, end_pad_width, data, lengths):
    pad_width = start_pad_width + end_pad_width
    out_size = data.shape[0] - (
        start_pad_width + end_pad_width) * len(lengths)
    out = np.ndarray((out_size,) + data.shape[1:])
    in_ptr = 0
    out_ptr = 0
    for length in lengths:
        out_length = length - pad_width
        out[out_ptr:(out_ptr + out_length)] = data[
            (in_ptr + start_pad_width):(in_ptr + length - end_pad_width)]
        in_ptr += length
        out_ptr += out_length
    lengths_out = lengths - (start_pad_width + end_pad_width)
    return (out, lengths_out)


def _gather_padding_ref(start_pad_width, end_pad_width, data, lengths):
    start_padding = np.zeros(data.shape[1:], dtype=data.dtype)
    end_padding = np.zeros(data.shape[1:], dtype=data.dtype)
    pad_width = start_pad_width + end_pad_width
    ptr = 0
    for length in lengths:
        for _ in range(start_pad_width):
            start_padding += data[ptr]
            ptr += 1
        ptr += length - pad_width
        for _ in range(end_pad_width):
            end_padding += data[ptr]
            ptr += 1
    return (start_padding, end_padding)


class TestSequenceOps(serial.SerializedTestCase):
    @given(start_pad_width=st.integers(min_value=1, max_value=2),
           end_pad_width=st.integers(min_value=0, max_value=2),
           args=_gen_test_add_padding(with_pad_data=True),
           ret_lengths=st.booleans(),
           **hu.gcs)
    @settings(deadline=10000)
    def test_add_padding(
        self, start_pad_width, end_pad_width, args, ret_lengths, gc, dc
    ):
        lengths, data, start_padding, end_padding = args
        start_padding = np.array(start_padding, dtype=np.float32)
        end_padding = np.array(end_padding, dtype=np.float32)
        outputs = ['output', 'lengths_out'] if ret_lengths else ['output']
        op = core.CreateOperator(
            'AddPadding', ['data', 'lengths', 'start_padding', 'end_padding'],
            outputs,
            padding_width=start_pad_width,
            end_padding_width=end_pad_width
        )
        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[data, lengths, start_padding, end_padding],
            reference=partial(
                _add_padding_ref, start_pad_width, end_pad_width, ret_lengths
            )
        )

    def _local_test_add_padding_shape_and_type(
        self,
        data,
        start_pad_width,
        end_pad_width,
        ret_lengths,
        lengths=None,
    ):
        if ret_lengths and lengths is None:
            return

        workspace.ResetWorkspace()
        workspace.FeedBlob("data", data)
        if lengths is not None:
            workspace.FeedBlob("lengths", np.array(lengths).astype(np.int32))

        op = core.CreateOperator(
            'AddPadding',
            ['data'] if lengths is None else ['data', 'lengths'],
            ['output', 'lengths_out'] if ret_lengths else ['output'],
            padding_width=start_pad_width,
            end_padding_width=end_pad_width
        )

        add_padding_net = core.Net("add_padding_net")
        add_padding_net.Proto().op.extend([op])
        assert workspace.RunNetOnce(
            add_padding_net
        ), "Failed to run the add_padding_net"

        shapes, types = workspace.InferShapesAndTypes(
            [add_padding_net],
        )

        expected_shape = list(data.shape)
        expected_shape[0] += (1 if lengths is None else len(lengths)) * (start_pad_width + end_pad_width)
        self.assertEqual(shapes["output"], expected_shape)
        self.assertEqual(types["output"], core.DataType.FLOAT)
        if ret_lengths:
            if lengths is None:
                self.assertEqual(shapes["lengths_out"], [1])
            else:
                self.assertEqual(shapes["lengths_out"], [len(lengths)])
            self.assertEqual(types["lengths_out"], core.DataType.INT32)


    def test_add_padding_shape_and_type_3(
        self
    ):
        for start_pad_width in range(3):
            for end_pad_width in range(3):
                for ret_lengths in [True, False]:
                    self._local_test_add_padding_shape_and_type(
                        data=np.random.rand(1, 2).astype(np.float32),
                        lengths=None,
                        start_pad_width=start_pad_width,
                        end_pad_width=end_pad_width,
                        ret_lengths=ret_lengths,
                    )

    def test_add_padding_shape_and_type_4(
        self
    ):
        for start_pad_width in range(3):
            for end_pad_width in range(3):
                for ret_lengths in [True, False]:
                    self._local_test_add_padding_shape_and_type(
                        data=np.random.rand(3, 1, 2).astype(np.float32),
                        lengths=[1, 1, 1],
                        start_pad_width=start_pad_width,
                        end_pad_width=end_pad_width,
                        ret_lengths=ret_lengths,
                    )

    def test_add_padding_shape_and_type_5(
        self
    ):
        for start_pad_width in range(3):
            for end_pad_width in range(3):
                for ret_lengths in [True, False]:
                    self._local_test_add_padding_shape_and_type(
                        data=np.random.rand(3, 2, 1).astype(np.float32),
                        lengths=None,
                        start_pad_width=start_pad_width,
                        end_pad_width=end_pad_width,
                        ret_lengths=ret_lengths,
                    )

    @given(start_pad_width=st.integers(min_value=0, max_value=3),
           end_pad_width=st.integers(min_value=0, max_value=3),
           num_dims=st.integers(min_value=1, max_value=4),
           num_groups=st.integers(min_value=0, max_value=4),
           ret_lengths=st.booleans(),
           **hu.gcs)
    @settings(deadline=1000)
    def test_add_padding_shape_and_type(
        self, start_pad_width, end_pad_width, num_dims, num_groups, ret_lengths, gc, dc
    ):
        np.random.seed(666)
        lengths = []
        for _ in range(num_groups):
            lengths.append(np.random.randint(0, 3))
        if sum(lengths) == 0:
            lengths = []

        data_shape = []
        for _ in range(num_dims):
            data_shape.append(np.random.randint(1, 4))
        if sum(lengths) > 0:
            data_shape[0] = sum(lengths)

        data = np.random.randn(*data_shape).astype(np.float32)

        self._local_test_add_padding_shape_and_type(
            data=data,
            lengths=lengths if len(lengths) else None,
            start_pad_width=start_pad_width,
            end_pad_width=end_pad_width,
            ret_lengths=ret_lengths,
        )

    @given(start_pad_width=st.integers(min_value=1, max_value=2),
           end_pad_width=st.integers(min_value=0, max_value=2),
           args=_gen_test_add_padding(with_pad_data=False),
           **hu.gcs)
    def test_add_zero_padding(self, start_pad_width, end_pad_width, args, gc, dc):
        lengths, data = args
        op = core.CreateOperator(
            'AddPadding',
            ['data', 'lengths'],
            ['output', 'lengths_out'],
            padding_width=start_pad_width,
            end_padding_width=end_pad_width)
        self.assertReferenceChecks(
            gc,
            op,
            [data, lengths],
            partial(_add_padding_ref, start_pad_width, end_pad_width, True))

    @given(start_pad_width=st.integers(min_value=1, max_value=2),
           end_pad_width=st.integers(min_value=0, max_value=2),
           data=hu.tensor(min_dim=1, max_dim=3),
           **hu.gcs)
    def test_add_padding_no_length(self, start_pad_width, end_pad_width, data, gc, dc):
        op = core.CreateOperator(
            'AddPadding',
            ['data'],
            ['output', 'output_lens'],
            padding_width=start_pad_width,
            end_padding_width=end_pad_width)
        self.assertReferenceChecks(
            gc,
            op,
            [data],
            partial(
                _add_padding_ref, start_pad_width, end_pad_width, True,
                lengths=np.array([data.shape[0]])))

    # Uncomment the following seed to make this fail.
    # @seed(302934307671667531413257853548643485645)
    # See https://github.com/caffe2/caffe2/issues/1547
    @unittest.skip("flaky test")
    @given(start_pad_width=st.integers(min_value=1, max_value=2),
           end_pad_width=st.integers(min_value=0, max_value=2),
           args=_gen_test_add_padding(with_pad_data=False, is_remove=True),
           **hu.gcs)
    def test_remove_padding(self, start_pad_width, end_pad_width, args, gc, dc):
        lengths, data = args
        op = core.CreateOperator(
            'RemovePadding',
            ['data', 'lengths'],
            ['output', 'lengths_out'],
            padding_width=start_pad_width,
            end_padding_width=end_pad_width)
        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[data, lengths],
            reference=partial(_remove_padding_ref, start_pad_width, end_pad_width))

    @given(start_pad_width=st.integers(min_value=0, max_value=2),
           end_pad_width=st.integers(min_value=0, max_value=2),
           args=_gen_test_add_padding(with_pad_data=True),
           **hu.gcs)
    @settings(deadline=10000)
    def test_gather_padding(self, start_pad_width, end_pad_width, args, gc, dc):
        lengths, data, start_padding, end_padding = args
        padded_data, padded_lengths = _add_padding_ref(
            start_pad_width, end_pad_width, True, data,
            lengths, start_padding, end_padding)
        op = core.CreateOperator(
            'GatherPadding',
            ['data', 'lengths'],
            ['start_padding', 'end_padding'],
            padding_width=start_pad_width,
            end_padding_width=end_pad_width)
        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[padded_data, padded_lengths],
            reference=partial(_gather_padding_ref, start_pad_width, end_pad_width))

    @given(data=hu.tensor(min_dim=3, max_dim=3, dtype=np.float32,
                          elements=hu.floats(min_value=-np.inf,
                                             max_value=np.inf),
                          min_value=1, max_value=10),
                          **hu.gcs)
    @settings(deadline=10000)
    def test_reverse_packed_segs(self, data, gc, dc):
        max_length = data.shape[0]
        batch_size = data.shape[1]
        lengths = np.random.randint(max_length + 1, size=batch_size)

        op = core.CreateOperator(
            "ReversePackedSegs",
            ["data", "lengths"],
            ["reversed_data"])

        def op_ref(data, lengths):
            rev_data = np.array(data, copy=True)
            for i in range(batch_size):
                seg_length = lengths[i]
                for j in range(seg_length):
                    rev_data[j][i] = data[seg_length - 1 - j][i]
            return (rev_data,)

        def op_grad_ref(grad_out, outputs, inputs):
            return op_ref(grad_out, inputs[1]) + (None,)

        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[data, lengths],
            reference=op_ref,
            output_to_grad='reversed_data',
            grad_reference=op_grad_ref)

    @given(data=hu.tensor(min_dim=1, max_dim=3, dtype=np.float32,
                          elements=hu.floats(min_value=-np.inf,
                                             max_value=np.inf),
                          min_value=10, max_value=10),
           indices=st.lists(st.integers(min_value=0, max_value=9),
                            min_size=0,
                            max_size=10),
           **hu.gcs_cpu_only)
    @settings(deadline=10000)
    def test_remove_data_blocks(self, data, indices, gc, dc):
        indices = np.array(indices)

        op = core.CreateOperator(
            "RemoveDataBlocks",
            ["data", "indices"],
            ["shrunk_data"])

        def op_ref(data, indices):
            unique_indices = np.unique(indices) if len(indices)>0 else np.array([],dtype=np.int64)
            sorted_indices = np.sort(unique_indices)
            shrunk_data = np.delete(data, sorted_indices, axis=0)
            return (shrunk_data,)

        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[data, indices],
            reference=op_ref)

    @given(elements=st.lists(st.integers(min_value=0, max_value=9),
                             min_size=0,
                             max_size=10),
           **hu.gcs_cpu_only)
    @settings(deadline=10000)
    def test_find_duplicate_elements(self, elements, gc, dc):
        mapping = {
            0: "a",
            1: "b",
            2: "c",
            3: "d",
            4: "e",
            5: "f",
            6: "g",
            7: "h",
            8: "i",
            9: "j"}
        data = np.array([mapping[e] for e in elements], dtype='|S')

        op = core.CreateOperator(
            "FindDuplicateElements",
            ["data"],
            ["indices"])

        def op_ref(data):
            unique_data = []
            indices = []
            for i, e in enumerate(data):
                if e in unique_data:
                    indices.append(i)
                else:
                    unique_data.append(e)
            return (np.array(indices, dtype=np.int64),)

        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[data],
            reference=op_ref)


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