from caffe2.python import core, workspace
from caffe2.python.test_util import TestCase
import numpy as np


class TestFeatureMapsOps(TestCase):

    def test_merge_dense_feature_tensors(self):
        op = core.CreateOperator(
            "MergeDenseFeatureTensors",
            [
                "in1", "in1_presence",
            ],
            [
                "out_lengths", "out_keys", "out_values",
            ],
            feature_ids=[11, 12, 13, 14]
        )
        # Input 1.
        workspace.FeedBlob(
            "in1",
            np.array([[11.1, 12.1, 13.1, 14.1], [11.2, 12.2, 13.2, 14.2]], dtype=np.float)
        )
        workspace.FeedBlob(
            "in1_presence",
            np.array([[True, False, False, True], [False, True, True, False]], dtype=np.bool)
        )

        workspace.RunOperatorOnce(op)

        np.testing.assert_array_equal(
            workspace.FetchBlob("out_lengths"),
            np.array([2, 2], dtype=np.int32)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_keys"),
            np.array([11, 14, 12, 13], dtype=np.int64)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_values"),
            np.array([11.1, 14.1, 12.2, 13.2], dtype=np.float)
        )


    def test_merge_single_scalar_feature_tensors(self):
        op = core.CreateOperator(
            "MergeSingleScalarFeatureTensors",
            [
                "in1", "in1_presence",
                "in2", "in2_presence",
            ],
            [
                "out_lengths", "out_keys", "out_values",
            ],
            feature_ids=[11, 12]
        )

        # Input 1.
        workspace.FeedBlob(
            "in1",
            np.array([11.1, 0.0], dtype=np.float)
        )
        workspace.FeedBlob(
            "in1_presence",
            np.array([True, False], dtype=np.bool)
        )
        # Input 2.
        workspace.FeedBlob(
            "in2",
            np.array([12.1, 12.2], dtype=np.float)
        )
        workspace.FeedBlob(
            "in2_presence",
            np.array([True, True], dtype=np.bool)
        )

        workspace.RunOperatorOnce(op)

        np.testing.assert_array_equal(
            workspace.FetchBlob("out_lengths"),
            np.array([2, 1], dtype=np.int32)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_keys"),
            np.array([11, 12, 12], dtype=np.int64)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_values"),
            np.array([11.1, 12.1, 12.2], dtype=np.float)
        )

    def test_merge_single_scalar_feature_tensors_gradient(self):
        op = core.CreateOperator(
            "MergeSingleScalarFeatureTensorsGradient",
            [
                "in1_presence",
                "in2_presence",
                "in3_presence",
                "out_values_grad",
            ],
            [
                "in1_grad", "in2_grad", "in3_grad",
            ],
        )

        # Inputs 1, 2 & 3.
        workspace.FeedBlob(
            "in1_presence",
            np.array([True, False], dtype=np.bool)
        )
        workspace.FeedBlob(
            "in2_presence",
            np.array([True, True], dtype=np.bool)
        )
        workspace.FeedBlob(
            "in3_presence",
            np.array([False, True], dtype=np.bool)
        )
        # Input 4.
        workspace.FeedBlob(
            "out_values_grad",
            np.array([0.1, 1.1, 1.2, 2.3], dtype=np.float)
        )

        workspace.RunOperatorOnce(op)

        np.testing.assert_array_equal(
            workspace.FetchBlob("in1_grad"),
            np.array([0.1, 0], dtype=np.float)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("in2_grad"),
            np.array([1.1, 1.2], dtype=np.float)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("in3_grad"),
            np.array([0, 2.3], dtype=np.float)
        )

    def test_merge_single_scalar_feature_tensors_gradient_with_strings(self):
        op = core.CreateOperator(
            "MergeSingleScalarFeatureTensorsGradient",
            [
                "in1_presence",
                "in2_presence",
                "in3_presence",
                "out_values_grad",
            ],
            [
                "in1_grad", "in2_grad", "in3_grad",
            ],
        )

        # Inputs 1, 2 & 3.
        workspace.FeedBlob(
            "in1_presence",
            np.array([True, False], dtype=np.bool)
        )
        workspace.FeedBlob(
            "in2_presence",
            np.array([True, True], dtype=np.bool)
        )
        workspace.FeedBlob(
            "in3_presence",
            np.array([False, True], dtype=np.bool)
        )
        # Input 4.
        workspace.FeedBlob(
            "out_values_grad",
            np.array(["0.1", "1.1", "1.2", "2.3"], dtype=np.unicode_)
        )

        workspace.RunOperatorOnce(op)

        np.testing.assert_array_equal(
            workspace.FetchBlob("in1_grad"),
            np.array(["0.1", ""], dtype=np.bytes_)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("in2_grad"),
            np.array(["1.1", "1.2"], dtype=np.bytes_)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("in3_grad"),
            np.array(["", "2.3"], dtype=np.bytes_)
        )

    def test_merge_single_list_feature_tensors(self):
        op = core.CreateOperator(
            "MergeSingleListFeatureTensors",
            [
                "in1_lengths", "in1_values", "in1_presence",
                "in2_lengths", "in2_values", "in2_presence",
            ],
            [
                "out_lengths", "out_keys", "out_values_lengths",
                "out_values_values",
            ],
            feature_ids=[11, 12]
        )

        # Input 1.
        workspace.FeedBlob(
            "in1_lengths",
            np.array([2, 0], dtype=np.int32)
        )
        workspace.FeedBlob(
            "in1_values",
            np.array([11.1, 11.2], dtype=np.float)
        )
        workspace.FeedBlob(
            "in1_presence",
            np.array([True, False], dtype=np.bool)
        )
        # Input 2.
        workspace.FeedBlob(
            "in2_lengths",
            np.array([2, 2], dtype=np.int32)
        )
        workspace.FeedBlob(
            "in2_values",
            np.array([12.1, 12.2, 12.3, 12.4], dtype=np.float)
        )
        workspace.FeedBlob(
            "in2_presence",
            np.array([True, True], dtype=np.bool)
        )

        workspace.RunOperatorOnce(op)

        np.testing.assert_array_equal(
            workspace.FetchBlob("out_lengths"),
            np.array([2, 1], dtype=np.int32)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_keys"),
            np.array([11, 12, 12], dtype=np.int64)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_values_lengths"),
            np.array([2, 2, 2], dtype=np.int32)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("out_values_values"),
            np.array([11.1, 11.2, 12.1, 12.2, 12.3, 12.4], dtype=np.float)
        )

    def test_merge_single_list_feature_tensors_gradient(self):
        self._test_merge_single_list_or_map_feature_tensors_gradient(
            "MergeSingleListFeatureTensorsGradient"
        )

    def test_merge_single_map_feature_tensors_gradient(self):
        self._test_merge_single_list_or_map_feature_tensors_gradient(
            "MergeSingleMapFeatureTensorsGradient"
        )

    def _test_merge_single_list_or_map_feature_tensors_gradient(self, op_name):
        op = core.CreateOperator(
            op_name,
            [
                "in1_lengths", "in1_presence",
                "in2_lengths", "in2_presence",
                "out_values_values_grad",
            ],
            [
                "in1_values_grad",
                "in2_values_grad",
            ],
        )

        # Input 1.
        workspace.FeedBlob(
            "in1_lengths",
            np.array([2, 0], dtype=np.int32)
        )
        workspace.FeedBlob(
            "in1_presence",
            np.array([True, False], dtype=np.bool)
        )
        # Input 2.
        workspace.FeedBlob(
            "in2_lengths",
            np.array([2, 2], dtype=np.int32)
        )
        workspace.FeedBlob(
            "in2_presence",
            np.array([True, True], dtype=np.bool)
        )
        workspace.FeedBlob(
            "out_values_values_grad",
            np.array([11.1, 11.2, 12.1, 12.2, 12.3, 12.4], dtype=np.float)
        )

        workspace.RunOperatorOnce(op)

        np.testing.assert_array_equal(
            workspace.FetchBlob("in1_values_grad"),
            np.array([11.1, 11.2], dtype=np.float)
        )
        np.testing.assert_array_equal(
            workspace.FetchBlob("in2_values_grad"),
            np.array([12.1, 12.2, 12.3, 12.4], dtype=np.float)
        )

    def test_merge_single_map_feature_tensors(self):
        op = core.CreateOperator(
            "MergeSingleMapFeatureTensors",
            [
                "in1_lengths", "in1_keys", "in1_values", "in1_presence",
                "in2_lengths", "in2_keys", "in2_values", "in2_presence",
            ],
            [
                "out_lengths", "out_keys", "out_values_lengths",
                "out_values_keys", "out_values_values",
            ],
            feature_ids=[11, 12]
        )

        # Input 1.
        workspace.FeedBlob(
            "in1_lengths",
            np.array([2, 0], dtype=np.int32)
        )
        workspace.FeedBlob(
            "in1_keys",
            np.array([111, 112], dtype=np.int64)
        )
        workspace.FeedBlob(
            "in1_values",
            np.array([11.1, 11.2], dtype=np.float)
        )
        workspace.FeedBlob(
            "in1_presence",
            np.array([True, False], dtype=np.bool)
        )
        # Input 2.
        workspace.FeedBlob(
            "in2_lengths",
            np.array([2, 2], dtype=np.int32)
        )