# Copyright (c) 2021 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
from op_test import (
    OpTest,
    convert_float_to_uint16,
    get_device_place,
    get_places,
    is_custom_device,
)
from scipy.special import psi

import paddle
from paddle import base, static
from paddle.base import core


class TestDigammaOp(OpTest):
    def setUp(self):
        # switch to static
        paddle.enable_static()

        self.op_type = 'digamma'
        self.python_api = paddle.digamma
        self.init_dtype_type()
        self.init_shape()
        data = np.random.random(self.shape).astype(self.dtype) + 1
        self.inputs = {'X': data}
        result = np.ones(self.shape).astype(self.dtype)
        result = psi(data)
        self.outputs = {'Out': result}

    def init_dtype_type(self):
        self.dtype = np.float64

    def init_shape(self):
        self.shape = (5, 32)

    def test_check_output(self):
        self.check_output(check_pir=True, check_symbol_infer=False)

    def test_check_grad_normal(self):
        self.check_grad(['X'], 'Out', check_pir=True)


class TestDigammaOpFp32(TestDigammaOp):
    def init_dtype_type(self):
        self.dtype = np.float32

    def test_check_grad_normal(self):
        self.check_grad(['X'], 'Out', check_pir=True)


class TestDigammaFP16Op(TestDigammaOp):
    def init_dtype_type(self):
        self.dtype = np.float16


class TestDigammaOp_ZeroSize(TestDigammaOp):
    def init_shape(self):
        self.shape = (5, 0)


@unittest.skipIf(
    not (core.is_compiled_with_cuda() or is_custom_device())
    or not core.is_bfloat16_supported(get_device_place()),
    "core is not compiled with CUDA or not support bfloat16",
)
class TestDigammaBF16Op(OpTest):
    def setUp(self):
        # switch to static
        paddle.enable_static()

        self.op_type = 'digamma'
        self.python_api = paddle.digamma
        self.init_dtype_type()
        shape = (5, 32)
        data = np.random.random(shape).astype(self.np_dtype) + 1
        self.inputs = {'X': convert_float_to_uint16(data)}
        result = np.ones(shape).astype(self.np_dtype)
        result = psi(data)
        self.outputs = {'Out': convert_float_to_uint16(result)}

    def init_dtype_type(self):
        self.dtype = np.uint16
        self.np_dtype = np.float32

    def test_check_output(self):
        # bfloat16 needs to set the parameter place
        self.check_output_with_place(
            get_device_place(), check_pir=True, check_symbol_infer=False
        )

    def test_check_grad_normal(self):
        self.check_grad_with_place(
            get_device_place(), ['X'], 'Out', check_pir=True
        )


class TestDigammaAPI(unittest.TestCase):
    def setUp(self):
        # switch to static
        paddle.enable_static()
        # prepare test attrs
        self.dtypes = ["float32", "float64"]
        self.places = get_places()
        self._shape = [8, 3, 32, 32]

    def test_in_static_mode(self):
        def init_input_output(dtype):
            input = np.random.random(self._shape).astype(dtype)
            return {'x': input}, psi(input)

        for dtype in self.dtypes:
            input_dict, sc_res = init_input_output(dtype)
            for place in self.places:
                with static.program_guard(static.Program()):
                    x = static.data(name="x", shape=self._shape, dtype=dtype)
                    out = paddle.digamma(x)

                    exe = static.Executor(place)
                    out_value = exe.run(feed=input_dict, fetch_list=[out])
                    np.testing.assert_allclose(out_value[0], sc_res, rtol=1e-05)

    def test_in_dynamic_mode(self):
        for dtype in self.dtypes:
            input = np.random.random(self._shape).astype(dtype)
            sc_res = psi(input)
            for place in self.places:
                # it is more convenient to use `guard` than `enable/disable_**` here
                with base.dygraph.guard(place):
                    input_t = paddle.to_tensor(input)
                    res = paddle.digamma(input_t).numpy()
                    np.testing.assert_allclose(res, sc_res, rtol=1e-05)

    def test_dtype_error(self):
        # in static graph mode
        with (
            self.assertRaises(TypeError),
            static.program_guard(static.Program()),
        ):
            x = static.data(name="x", shape=self._shape, dtype="bool")
            out = paddle.digamma(x, name="digamma_res")

        # in dynamic mode
        with (
            self.assertRaises(RuntimeError),
            base.dygraph.guard(),
        ):
            input = np.random.random(self._shape).astype("bool")
            input_t = paddle.to_tensor(input)
            res = paddle.digamma(input_t)


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