# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# 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.
import sys
import unittest
import uuid

sys.path.insert(0, '.')

import config
import numpy as np
import utils
from numpy.random import randint, randn

import paddle
from paddle.incubate.autograd import primops

paddle.enable_static()


@utils.place(config.DEVICES)
@utils.parameterize(
    (
        utils.TEST_CASE_NAME,
        'op',
        'args',
        'kwargs',
        'expected_shape',
        'expected_dtype',
    ),
    (
        ('add', primops.add, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
        ('sub', primops.sub, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
        ('mul', primops.mul, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
        ('div', primops.div, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
        ('sub', primops.sub, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
        ('sqrt', primops.sqrt, randn(2, 3), {}, (2, 3), 'float64'),
        ('tanh', primops.tanh, randn(2, 3), {}, (2, 3), 'float64'),
        ('sin', primops.sin, randn(2, 3), {}, (2, 3), 'float64'),
        ('cos', primops.cos, randn(2, 3), {}, (2, 3), 'float64'),
        ('exp', primops.exp, randn(2, 3), {}, (2, 3), 'float64'),
        ('erf', primops.erf, randn(2, 3), {}, (2, 3), 'float64'),
        ('abs', primops.abs, randn(2, 3), {}, (2, 3), 'float64'),
        ('log', primops.log, randn(2, 3), {}, (2, 3), 'float64'),
        (
            'cast',
            primops.cast,
            randn(2, 3),
            {'dtype': paddle.int64},
            (2, 3),
            'int64',
        ),
        (
            'reshape',
            primops.reshape,
            randn(2, 3),
            {'shape': (3, 2)},
            (3, 2),
            'float64',
        ),
        (
            'broadcast',
            primops.broadcast,
            randn(2),
            {'shape': (3, 2)},
            (3, 2),
            'float64',
        ),
        (
            'transpose',
            primops.transpose,
            randn(2, 3),
            {'axis': (1, 0)},
            (3, 2),
            'float64',
        ),
        (
            'concat_axis0',
            primops.concat,
            ((randn(2, 3), randn(2, 3)),),
            {'axis': 0},
            (4, 3),
            'float64',
        ),
        (
            'concat_axis1',
            primops.concat,
            ((randn(2, 3), randn(2, 3)),),
            {'axis': 1},
            (2, 6),
            'float64',
        ),
        (
            'reduce_axis1',
            primops.reduce_sum,
            randn(2, 3),
            {'axis': (1,)},
            (2,),
            'float64',
        ),
        (
            'reduce_axis01',
            primops.reduce_sum,
            randn(2, 3),
            {'axis': (0, 1)},
            (),
            'float64',
        ),
        (
            'split',
            primops.split,
            randn(2, 3),
            {'num_or_sections': [1, 2], 'axis': 1},
            ((2, 1), (2, 2)),
            ('float64', 'float64'),
        ),
        (
            'matmul',
            primops.matmul,
            (randn(2, 3), randn(3, 2)),
            {},
            (2, 2),
            'float64',
        ),
        (
            'slice_select',
            primops.slice_select,
            randn(3, 2),
            {'axis': [0], 'starts': [0], 'ends': [2], 'strides': [1]},
            (2, 2),
            'float64',
        ),
        (
            'slice_assign',
            primops.slice_assign,
            (randn(2, 3), randn(2, 2)),
            {'axis': [1], 'starts': [1], 'ends': [3], 'strides': [1]},
            (2, 3),
            'float64',
        ),
        (
            'gather',
            primops.gather,
            (randn(3, 2), randint(0, 2, (5,), np.int32)),
            {'axis': 0},
            (5, 2),
            'float64',
        ),
        (
            'scatter_add',
            primops.scatter_add,
            (randn(3, 2), randn(5, 2), randint(0, 2, (5,), np.int32)),
            {'axis': 0},
            (3, 2),
            'float64',
        ),
        (
            'fill_const',
            primops.fill_const,
            (),
            {'value': 10, 'shape': (3, 2), 'dtype': paddle.float32},
            (3, 2),
            'float32',
        ),
        ('neg', primops.neg, randn(2, 3), {}, (2, 3), 'float64'),
        (
            'select',
            primops.select,
            (randn(2, 3) > 0, randn(2, 3), randn(2, 3)),
            {},
            (2, 3),
            'float64',
        ),
        ('eq', primops.eq, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'bool'),
        ('ne', primops.ne, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'bool'),
        ('gt', primops.gt, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'bool'),
        ('ge', primops.ge, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'bool'),
        ('pow', primops.pow, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
        ('max', primops.max, (randn(2, 3), randn(2, 3)), {}, (2, 3), 'float64'),
    ),
)
class TestPrimops(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        paddle.enable_static()

    @classmethod
    def tearDownClass(cls):
        paddle.disable_static()

    def test_prim_ops(self):
        program = paddle.static.Program()
        with paddle.static.program_guard(program):
            args = self._as_tuple(self.args)
            args = self.arr2var(args)
            results = self.op(*args, **self.kwargs)
            results = self._as_tuple(results)
            expected_shape = self._as_tuple(self.expected_shape)
            expected_dtype = self._as_tuple(self.expected_dtype)

            for r, shape, dtype in zip(results, expected_shape, expected_dtype):
                self.assertEqual(r.shape, shape)
                self.assertEqual(str(r.dtype).split('.')[1], dtype)

    def arr2var(self, arr):
        """convert numpy ndarray to paddle Variable recursively."""
        return [
            paddle.static.data(f'x{uuid.uuid4()}', v.shape, v.dtype)
            if isinstance(v, np.ndarray)
            else self.arr2var(v)
            for v in arr
        ]

    def _as_tuple(self, input):
        if isinstance(input, (tuple, list)) and len(input) == 0:
            return input
        if not isinstance(input, (tuple, list)) or all(
            isinstance(i, int) for i in input
        ):
            return (input,)
        return input


if __name__ == '__main__':
    unittest.main()
