# 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 unittest

import numpy as np

import paddle
import paddle.fluid as fluid
from paddle.fluid.framework import _test_eager_guard


def numpy_cov(np_arr, rowvar=True, ddof=1, fweights=None, aweights=None):
    return np.cov(
        np_arr,
        rowvar=rowvar,
        ddof=int(ddof),
        fweights=fweights,
        aweights=aweights,
    )


class Cov_Test(unittest.TestCase):
    def setUp(self):
        self.shape = [20, 10]
        self.weightshape = [10]

    def func_test_tensor_cov_default(self):
        typelist = ['float64']
        places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            places.append(fluid.CUDAPlace(0))

        for idx, p in enumerate(places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')

            for dtype in typelist:
                np_arr = np.random.rand(*self.shape).astype(dtype)
                tensor = paddle.to_tensor(np_arr, place=p)
                cov = paddle.linalg.cov(
                    tensor, rowvar=True, ddof=True, fweights=None, aweights=None
                )
                np_cov = numpy_cov(
                    np_arr, rowvar=True, ddof=1, fweights=None, aweights=None
                )
                np.testing.assert_allclose(np_cov, cov.numpy(), rtol=1e-05)

    def test_tensor_cov_default(self):
        with _test_eager_guard():
            self.func_test_tensor_cov_default()
        self.func_test_tensor_cov_default()

    def func_test_tensor_cov_rowvar(self):
        typelist = ['float64']
        places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            places.append(fluid.CUDAPlace(0))

        for idx, p in enumerate(places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')

            for dtype in typelist:
                np_arr = np.random.rand(*self.shape).astype(dtype)
                tensor = paddle.to_tensor(np_arr, place=p)
                cov = paddle.linalg.cov(
                    tensor,
                    rowvar=False,
                    ddof=True,
                    fweights=None,
                    aweights=None,
                )
                np_cov = numpy_cov(
                    np_arr, rowvar=False, ddof=1, fweights=None, aweights=None
                )
                np.testing.assert_allclose(np_cov, cov.numpy(), rtol=1e-05)

    def test_tensor_cov_rowvar(self):
        with _test_eager_guard():
            self.func_test_tensor_cov_rowvar()
        self.func_test_tensor_cov_rowvar()

    def func_test_tensor_cov_ddof(self):
        typelist = ['float64']
        places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            places.append(fluid.CUDAPlace(0))

        for idx, p in enumerate(places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')

            for dtype in typelist:
                np_arr = np.random.rand(*self.shape).astype(dtype)
                tensor = paddle.to_tensor(np_arr, place=p)
                cov = paddle.linalg.cov(
                    tensor,
                    rowvar=True,
                    ddof=False,
                    fweights=None,
                    aweights=None,
                )
                np_cov = numpy_cov(
                    np_arr, rowvar=True, ddof=0, fweights=None, aweights=None
                )
                np.testing.assert_allclose(np_cov, cov.numpy(), rtol=1e-05)

    def test_tensor_cov_ddof(self):
        with _test_eager_guard():
            self.func_test_tensor_cov_ddof()
        self.func_test_tensor_cov_ddof()

    def func_test_tensor_cov_fweights(self):
        typelist = ['float64']
        places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            places.append(fluid.CUDAPlace(0))

        for idx, p in enumerate(places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')

            for dtype in typelist:
                np_arr = np.random.rand(*self.shape).astype(dtype)
                np_fw = np.random.randint(10, size=self.weightshape).astype(
                    'int32'
                )
                tensor = paddle.to_tensor(np_arr, place=p)
                fweights = paddle.to_tensor(np_fw, place=p)
                cov = paddle.linalg.cov(
                    tensor,
                    rowvar=True,
                    ddof=True,
                    fweights=fweights,
                    aweights=None,
                )
                np_cov = numpy_cov(
                    np_arr, rowvar=True, ddof=1, fweights=np_fw, aweights=None
                )
                np.testing.assert_allclose(np_cov, cov.numpy(), rtol=1e-05)

    def test_tensor_cov_fweights(self):
        with _test_eager_guard():
            self.func_test_tensor_cov_fweights()
        self.func_test_tensor_cov_fweights()

    def func_test_tensor_cov_aweights(self):
        typelist = ['float64']
        places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            places.append(fluid.CUDAPlace(0))

        for idx, p in enumerate(places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')

            for dtype in typelist:
                np_arr = np.random.rand(*self.shape).astype(dtype)
                np_aw = np.random.randint(10, size=self.weightshape).astype(
                    'int32'
                )
                tensor = paddle.to_tensor(np_arr, place=p)
                aweights = paddle.to_tensor(np_aw, place=p)
                cov = paddle.linalg.cov(
                    tensor,
                    rowvar=True,
                    ddof=True,
                    fweights=None,
                    aweights=aweights,
                )
                np_cov = numpy_cov(
                    np_arr, rowvar=True, ddof=1, fweights=None, aweights=np_aw
                )
                np.testing.assert_allclose(np_cov, cov.numpy(), rtol=1e-05)

    def test_tensor_cov_aweights(self):
        with _test_eager_guard():
            self.func_test_tensor_cov_aweights()
        self.func_test_tensor_cov_aweights()

    def func_test_tensor_cov_weights(self):
        typelist = ['float64']
        places = [fluid.CPUPlace()]
        if fluid.core.is_compiled_with_cuda():
            places.append(fluid.CUDAPlace(0))

        for idx, p in enumerate(places):
            if idx == 0:
                paddle.set_device('cpu')
            else:
                paddle.set_device('gpu')

            for dtype in typelist:
                np_arr = np.random.rand(*self.shape).astype(dtype)
                np_fw = np.random.randint(10, size=self.weightshape).astype(
                    'int64'
                )
                np_aw = np.random.rand(*self.weightshape).astype('float64')
                tensor = paddle.to_tensor(np_arr, place=p)
                fweights = paddle.to_tensor(np_fw, place=p)
                aweights = paddle.to_tensor(np_aw, place=p)
                cov = paddle.linalg.cov(
                    tensor,
                    rowvar=True,
                    ddof=True,
                    fweights=fweights,
                    aweights=aweights,
                )
                np_cov = numpy_cov(
                    np_arr, rowvar=True, ddof=1, fweights=np_fw, aweights=np_aw
                )
                np.testing.assert_allclose(np_cov, cov.numpy(), rtol=1e-05)

    def test_tensor_cov_weights(self):
        with _test_eager_guard():
            self.func_test_tensor_cov_weights()
        self.func_test_tensor_cov_weights()


class Cov_Test2(Cov_Test):
    def setUp(self):
        self.shape = [10]
        self.weightshape = [10]


# Input(x) only support N-D (1<=N<=2) tensor
class Cov_Test3(unittest.TestCase):
    def setUp(self):
        self.shape = [2, 5, 10]
        self.fweightshape = [10]
        self.aweightshape = [10]
        self.fw_s = 1.0
        self.aw_s = 1.0

    def func_test_errors(self):
        def test_err():
            np_arr = np.random.rand(*self.shape).astype('float64')
            np_fw = self.fw_s * np.random.rand(*self.fweightshape).astype(
                'int32'
            )
            np_aw = self.aw_s * np.random.rand(*self.aweightshape).astype(
                'float64'
            )
            tensor = paddle.to_tensor(np_arr)
            fweights = paddle.to_tensor(np_fw)
            aweights = paddle.to_tensor(np_aw)
            cov = paddle.linalg.cov(
                tensor,
                rowvar=True,
                ddof=True,
                fweights=fweights,
                aweights=aweights,
            )

        self.assertRaises(ValueError, test_err)

    def test_errors(self):
        with _test_eager_guard():
            self.func_test_errors()
        self.func_test_errors()


# Input(fweights) only support N-D (N<=1) tensor
class Cov_Test4(Cov_Test3):
    def setUp(self):
        self.shape = [5, 10]
        self.fweightshape = [2, 10]
        self.aweightshape = [10]
        self.fw_s = 1.0
        self.aw_s = 1.0


# The number of Input(fweights) should equal to x's dim[1]
class Cov_Test5(Cov_Test3):
    def setUp(self):
        self.shape = [5, 10]
        self.fweightshape = [5]
        self.aweightshape = [10]
        self.fw_s = 1.0
        self.aw_s = 1.0


# The value of Input(fweights) cannot be negtive
class Cov_Test6(Cov_Test3):
    def setUp(self):
        self.shape = [5, 10]
        self.fweightshape = [10]
        self.aweightshape = [10]
        self.fw_s = -1.0
        self.aw_s = 1.0


# Input(aweights) only support N-D (N<=1) tensor
class Cov_Test7(Cov_Test3):
    def setUp(self):
        self.shape = [5, 10]
        self.fweightshape = [10]
        self.aweightshape = [2, 10]
        self.fw_s = 1.0
        self.aw_s = 1.0


# The number of Input(aweights) should equal to x's dim[1]
class Cov_Test8(Cov_Test3):
    def setUp(self):
        self.shape = [5, 10]
        self.fweightshape = [10]
        self.aweightshape = [5]
        self.fw_s = 1.0
        self.aw_s = 1.0


# The value of Input(aweights) cannot be negtive
class Cov_Test9(Cov_Test3):
    def setUp(self):
        self.shape = [5, 10]
        self.fweightshape = [10]
        self.aweightshape = [10]
        self.fw_s = 1.0
        self.aw_s = -1.0


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