# Copyright (c) 2016-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

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 caffe2.python.test_util import TestCase, rand_array


class TestPartitionOps(TestCase):

    def test_configs(self):
        # (main dims, partitions,  main type, [list of (extra dims, type)])
        configs = [
            ((10, ), 3),
            ((4, ), 10),
            ((10, 10), 4),
            ((100, ), 2),
            ((5, ), 1),
            ((1, ), 1),
            ((2, 10), 2),
        ]
        suffixes = [
            [],
            [((2, 2), np.float32)],
            [((3, ), np.int64), ((2, ), np.float32)],
        ]
        return [
            (main_dims, parts, main_type, extra, pack)
            for main_dims, parts in configs
            for main_type in [np.int32, np.int64] for extra in suffixes
            for pack in [False, True]
        ]

    def testPartition(self):
        for main_dims, parts, main_type, extra_ins, pack in self.test_configs():
            ins = ['in' + str(i) for i in range(1 + len(extra_ins))]
            outs = [
                'in{}_p{}'.format(j, i)
                for i in range(parts) for j in range(1 + len(extra_ins))
            ]
            op = core.CreateOperator(
                'Partition', ins, outs, pack_first_input=(1 if pack else 0))
            x = []
            for i, (dims, t) in enumerate([((), main_type)] + extra_ins):
                if t in [np.float32, np.float64]:
                    d = rand_array(*(main_dims + dims))
                else:
                    d = np.random.randint(-100, 100, (main_dims + dims))
                d = d.astype(t)
                workspace.FeedBlob(ins[i], d)
                x.append(d)

            def sharding(x):
                # numpy has proper modulo op that yields non-negative results
                shards = (x[0] % parts).reshape([-1])
                out = []
                for i in range(parts):
                    for ind, v in enumerate(x):
                        suffix_shape = v.shape[len(x[0].shape):]
                        accum = []
                        data = v.reshape((-1, ) + suffix_shape)

                        if pack and ind == 0:
                            data = data // parts

                        for j, s in enumerate(shards):
                            if s == i:
                                accum.append(data[j])

                        def join(a):
                            if not a:
                                return np.empty(shape=(0, ) + suffix_shape)
                            return np.stack(a)

                        out.append(join(accum))
                return out

            workspace.RunOperatorOnce(op)
            ref = sharding(x)
            print(x)
            print(ref)
            for name, expected in zip(outs, ref):
                np.testing.assert_array_equal(
                    expected, workspace.FetchBlob(name)
                )

            # test inverse operation (GatherByKey)
            if len(main_dims) == 1:
                # currently only 1D key tensor supported
                for i in range(len(extra_ins)):
                    expected_out = ins[i + 1]
                    gather_ins = [ins[0]] + [
                        outs[len(ins) * p + i + 1] for p in range(parts)]
                    actual_out = expected_out + '_actual'
                    op = core.CreateOperator(
                        'GatherByKey', gather_ins, actual_out)
                    workspace.RunOperatorOnce(op)
                    expected = workspace.FetchBlob(expected_out)
                    actual = workspace.FetchBlob(actual_out)
                    np.testing.assert_array_equal(expected, actual)


    def testLengthsPartition(self):
        for main_dims, parts, main_type, extra_ins, pack in self.test_configs():
            # For LengthsSharding only 1-D tensors supported as a first input
            if len(main_dims) > 1:
                continue
            ins = ['in' + str(i) for i in range(2 + len(extra_ins))]
            outs = [
                'in{}_p{}'.format(j, i)
                for i in range(parts) for j in range(2 + len(extra_ins))
            ]
            op = core.CreateOperator(
                'LengthsPartition', ins, outs,
                pack_first_input=(1 if pack else 0)
            )
            x = []
            for i, (dims, t) in enumerate([((), main_type)] + extra_ins):
                if t in [np.float32, np.float64]:
                    d = rand_array(*(main_dims + dims))
                else:
                    d = np.random.randint(-100, 100, (main_dims + dims))
                d = d.astype(t)
                workspace.FeedBlob(ins[i + 1], d)
                x.append(d)

            # Randomly generate length tensor as well
            elements = np.random.randint(2, 10)
            lengths = []
            total_length = 0
            for _ in range(elements - 1):
                lengths.append(np.random.randint(main_dims[0] - total_length))
                total_length += lengths[-1]
            lengths.append(main_dims[0] - total_length)
            workspace.FeedBlob(ins[0], np.array(lengths, dtype=np.int32))

            def sharding(x):
                # numpy has proper modulo op that yields non-negative results
                shards = (x[0] % parts).reshape([-1])
                out = []
                for i in range(parts):
                    idx = 0
                    sharded_lengths = np.zeros(elements)
                    for ind, length in enumerate(lengths):
                        for _ in range(length):
                            if shards[idx] == i:
                                sharded_lengths[ind] += 1
                            idx += 1
                    out.append(sharded_lengths)

                    for ind, v in enumerate(x):
                        suffix_shape = v.shape[len(x[0].shape):]
                        accum = []
                        data = v.reshape((-1, ) + suffix_shape)

                        if pack and ind == 0:
                            data = data // parts

                        for j, s in enumerate(shards):
                            if s == i:
                                accum.append(data[j])

                        def join(a):
                            if not a:
                                return np.empty(shape=(0, ) + suffix_shape)
                            return np.stack(a)

                        out.append(join(accum))
                return out

            workspace.RunOperatorOnce(op)
            ref = sharding(x)
            for name, expected in zip(outs, ref):
                np.testing.assert_array_equal(
                    expected, workspace.FetchBlob(name)
                )

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

[((196, 4, 196, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((57, 17, 58, 76), 'caffe2.python.core.CreateOperator', 'core.CreateOperator', (), '', False, 'from caffe2.python import core, workspace\n'), ((94, 12, 94, 41), 'caffe2.python.workspace.RunOperatorOnce', 'workspace.RunOperatorOnce', ({(94, 38, 94, 40): 'op'}, {}), '(op)', False, 'from caffe2.python import core, workspace\n'), ((129, 17, 132, 13), 'caffe2.python.core.CreateOperator', 'core.CreateOperator', (), '', False, 'from caffe2.python import core, workspace\n'), ((144, 23, 144, 47), 'numpy.random.randint', 'np.random.randint', ({(144, 41, 144, 42): '2', (144, 44, 144, 46): '10'}, {}), '(2, 10)', True, 'import numpy as np\n'), ((187, 12, 187, 41), 'caffe2.python.workspace.RunOperatorOnce', 'workspace.RunOperatorOnce', ({(187, 38, 187, 40): 'op'}, {}), '(op)', False, 'from caffe2.python import core, workspace\n'), ((66, 16, 66, 45), 'caffe2.python.workspace.FeedBlob', 'workspace.FeedBlob', ({(66, 35, 66, 41): 'ins[i]', (66, 43, 66, 44): 'd'}, {}), '(ins[i], d)', False, 'from caffe2.python import core, workspace\n'), ((140, 16, 140, 49), 'caffe2.python.workspace.FeedBlob', 'workspace.FeedBlob', ({(140, 35, 140, 45): 'ins[i + 1]', (140, 47, 140, 48): 'd'}, {}), '(ins[i + 1], d)', False, 'from caffe2.python import core, workspace\n'), ((151, 39, 151, 72), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((62, 24, 62, 55), 'caffe2.python.test_util.rand_array', 'rand_array', ({(62, 35, 62, 53): '*(main_dims + dims)'}, {}), '(*(main_dims + dims))', False, 'from caffe2.python.test_util import TestCase, rand_array\n'), ((64, 24, 64, 72), 'numpy.random.randint', 'np.random.randint', ({(64, 42, 64, 46): '-100', (64, 48, 64, 51): '100', (64, 54, 64, 70): 'main_dims + dims'}, {}), '(-100, 100, main_dims + dims)', True, 'import numpy as np\n'), ((100, 30, 100, 55), 'caffe2.python.workspace.FetchBlob', 'workspace.FetchBlob', ({(100, 50, 100, 54): 'name'}, {}), '(name)', False, 'from caffe2.python import core, workspace\n'), ((111, 25, 112, 62), 'caffe2.python.core.CreateOperator', 'core.CreateOperator', ({(112, 24, 112, 37): '"""GatherByKey"""', (112, 39, 112, 49): 'gather_ins', (112, 51, 112, 61): 'actual_out'}, {}), "('GatherByKey', gather_ins, actual_out)", False, 'from caffe2.python import core, workspace\n'), ((113, 20, 113, 49), 'caffe2.python.workspace.RunOperatorOnce', 'workspace.RunOperatorOnce', ({(113, 46, 113, 48): 'op'}, {}), '(op)', False, 'from caffe2.python import core, workspace\n'), ((114, 31, 114, 64), 'caffe2.python.workspace.FetchBlob', 'workspace.FetchBlob', ({(114, 51, 114, 63): 'expected_out'}, {}), '(expected_out)', False, 'from caffe2.python import core, workspace\n'), ((115, 29, 115, 60), 'caffe2.python.workspace.FetchBlob', 'workspace.FetchBlob', ({(115, 49, 115, 59): 'actual_out'}, {}), '(actual_out)', False, 'from caffe2.python import core, workspace\n'), ((116, 20, 116, 67), 'numpy.testing.assert_array_equal', 'np.testing.assert_array_equal', ({(116, 50, 116, 58): 'expected', (116, 60, 116, 66): 'actual'}, {}), '(expected, actual)', True, 'import numpy as np\n'), ((136, 24, 136, 55), 'caffe2.python.test_util.rand_array', 'rand_array', ({(136, 35, 136, 53): '*(main_dims + dims)'}, {}), '(*(main_dims + dims))', False, 'from caffe2.python.test_util import TestCase, rand_array\n'), ((138, 24, 138, 72), 'numpy.random.randint', 'np.random.randint', ({(138, 42, 138, 46): '-100', (138, 48, 138, 51): '100', (138, 54, 138, 70): 'main_dims + dims'}, {}), '(-100, 100, main_dims + dims)', True, 'import numpy as np\n'), ((148, 31, 148, 77), 'numpy.random.randint', 'np.random.randint', ({(148, 49, 148, 76): '(main_dims[0] - total_length)'}, {}), '(main_dims[0] - total_length)', True, 'import numpy as np\n'), ((159, 38, 159, 56), 'numpy.zeros', 'np.zeros', ({(159, 47, 159, 55): 'elements'}, {}), '(elements)', True, 'import numpy as np\n'), ((191, 30, 191, 55), 'caffe2.python.workspace.FetchBlob', 'workspace.FetchBlob', ({(191, 50, 191, 54): 'name'}, {}), '(name)', False, 'from caffe2.python import core, workspace\n'), ((89, 35, 89, 46), 'numpy.stack', 'np.stack', ({(89, 44, 89, 45): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((182, 35, 182, 46), 'numpy.stack', 'np.stack', ({(182, 44, 182, 45): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((88, 39, 88, 75), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((181, 39, 181, 75), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n')]