# Copyright (c) 2019 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 paddle.fluid as fluid
import paddle.nn.functional as F
import unittest
import numpy as np
import six
import paddle
from paddle.fluid.framework import _test_eager_guard


class TensorFillDiagTensor_Test(unittest.TestCase):

    def setUp(self):
        self.typelist = ['float32', 'float64', 'int32', 'int64']
        self.places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            self.places.append(fluid.CUDAPlace(0))

    def func_dim2(self):
        expected_np = np.array([[1, 2, 2], [2, 1, 2], [2, 2, 1],
                                [2, 2, 2]]).astype('float32')
        expected_grad = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0],
                                  [1, 1, 1]]).astype('float32')

        for idx, p in enumerate(self.places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')
            for dtype in self.typelist:
                v = paddle.ones((3, ), dtype=dtype)
                var = (np.random.random() + 1)
                x = paddle.ones((4, 3), dtype=dtype)
                x.stop_gradient = False
                y = x * 2
                y.fill_diagonal_tensor_(v, offset=0, dim1=0, dim2=1)
                loss = y.sum()
                loss.backward()

                self.assertEqual(
                    (y.numpy().astype('float32') == expected_np).all(), True)
                self.assertEqual(
                    (y.grad.numpy().astype('float32') == expected_grad).all(),
                    True)

    def test_dim2(self):
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
        with _test_eager_guard():
            self.func_dim2()
        self.func_dim2()
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": False})

    def func_dim2_offset_1(self):
        expected_np = np.array([[2, 2, 2], [1, 2, 2], [2, 1, 2],
                                [2, 2, 1]]).astype('float32')
        expected_grad = np.array([[1, 1, 1], [0, 1, 1], [1, 0, 1],
                                  [1, 1, 0]]).astype('float32')

        for idx, p in enumerate(self.places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')
            for dtype in self.typelist:
                v = paddle.ones((3, ), dtype=dtype)
                var = (np.random.random() + 1)
                x = paddle.ones((4, 3), dtype=dtype)
                x.stop_gradient = False
                y = x * 2
                y.fill_diagonal_tensor_(v, offset=-1, dim1=0, dim2=1)
                loss = y.sum()
                loss.backward()

                self.assertEqual(
                    (y.numpy().astype('float32') == expected_np).all(), True)
                self.assertEqual(
                    (y.grad.numpy().astype('float32') == expected_grad).all(),
                    True)

    def test_dim2_offset_1(self):
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
        with _test_eager_guard():
            self.func_dim2_offset_1()
        self.func_dim2_offset_1()
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": False})

    def func_dim2_offset1(self):
        expected_np = np.array([[2, 1, 2], [2, 2, 1], [2, 2, 2],
                                [2, 2, 2]]).astype('float32')
        expected_grad = np.array([[1, 0, 1], [1, 1, 0], [1, 1, 1],
                                  [1, 1, 1]]).astype('float32')

        for idx, p in enumerate(self.places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')
            for dtype in self.typelist:
                v = paddle.ones((2, ), dtype=dtype)
                var = (np.random.random() + 1)
                x = paddle.ones((4, 3), dtype=dtype)
                x.stop_gradient = False
                y = x * 2
                y.fill_diagonal_tensor_(v, offset=1, dim1=0, dim2=1)
                loss = y.sum()
                loss.backward()

                self.assertEqual(
                    (y.numpy().astype('float32') == expected_np).all(), True)
                self.assertEqual(
                    (y.grad.numpy().astype('float32') == expected_grad).all(),
                    True)

    def test_dim2_offset1(self):
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
        with _test_eager_guard():
            self.func_dim2_offset1()
        self.func_dim2_offset1()
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": False})

    def func_dim4(self):
        expected_np = np.array([[[[0, 3], [2, 2], [2, 2]],
                                 [[2, 2], [1, 4], [2, 2]],
                                 [[2, 2], [2, 2], [2, 5]],
                                 [[2, 2], [2, 2], [2, 2]]],
                                [[[6, 9], [2, 2], [2, 2]],
                                 [[2, 2], [7, 10], [2, 2]],
                                 [[2, 2], [2, 2], [8, 11]],
                                 [[2, 2], [2, 2], [2, 2]]]]).astype('float32')
        expected_grad = np.array([[[[0, 0], [1, 1], [1, 1]],
                                   [[1, 1], [0, 0], [1, 1]],
                                   [[1, 1], [1, 1], [0, 0]],
                                   [[1, 1], [1, 1], [1, 1]]],
                                  [[[0, 0], [1, 1], [1, 1]],
                                   [[1, 1], [0, 0], [1, 1]],
                                   [[1, 1], [1, 1], [0, 0]],
                                   [[1, 1], [1, 1], [1, 1]]]]).astype('float32')

        for idx, p in enumerate(self.places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')
            for dtype in self.typelist:
                v = paddle.to_tensor(np.arange(12).reshape(2, 2, 3),
                                     dtype=dtype)
                var = (np.random.random() + 1)
                x = paddle.ones((2, 4, 3, 2), dtype=dtype)
                x.stop_gradient = False
                y = x * 2
                y.fill_diagonal_tensor_(v, offset=0, dim1=1, dim2=2)
                loss = y.sum()
                loss.backward()

                self.assertEqual(
                    (y.numpy().astype('float32') == expected_np).all(), True)
                self.assertEqual(
                    (y.grad.numpy().astype('float32') == expected_grad).all(),
                    True)

    def test_func_dim4(self):
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
        with _test_eager_guard():
            self.func_dim4()
        self.func_dim4()
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": False})

    def func_largedim(self):
        #large dim only test on gpu because the cpu version is too slow for ci test, and the memory is limited
        if len(self.places) > 1:
            bsdim = 1024
            fsdim = 128
            paddle.set_device('gpu')
            for dtype in self.typelist:
                v = paddle.arange(bsdim * fsdim, dtype=dtype).reshape(
                    (bsdim, fsdim))
                y = paddle.ones((bsdim, fsdim, fsdim), dtype=dtype)
                y.stop_gradient = False
                y = y * 2
                y.fill_diagonal_tensor_(v, offset=0, dim1=1, dim2=2)
                loss = y.sum()
                loss.backward()

                expected_pred = v - 2
                expected_pred = F.diag_embed(expected_pred) + 2
                expected_grad = paddle.ones(v.shape, dtype=dtype) - 2
                expected_grad = F.diag_embed(expected_grad) + 1

                self.assertEqual((y == expected_pred).all(), True)
                self.assertEqual((y.grad == expected_grad).all(), True)

    def test_largedim(self):
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
        with _test_eager_guard():
            self.func_largedim()
        self.func_largedim()
        fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": False})


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