# 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
import paddle

import config
import mock_data as mock
import parameterize

np.random.seed(2022)


@parameterize.place(config.DEVICES)
@parameterize.parameterize_cls(
    (parameterize.TEST_CASE_NAME, 'dist'),
    [('test-mock-exp',
      mock.Exponential(
          rate=paddle.rand([100, 200, 99], dtype=config.DEFAULT_DTYPE)))])
class TestExponentialFamily(unittest.TestCase):

    def test_entropy(self):
        np.testing.assert_allclose(
            self.dist.entropy(),
            paddle.distribution.ExponentialFamily.entropy(self.dist),
            rtol=config.RTOL.get(config.DEFAULT_DTYPE),
            atol=config.ATOL.get(config.DEFAULT_DTYPE))


@parameterize.place(config.DEVICES)
@parameterize.parameterize_cls(
    (config.TEST_CASE_NAME, 'dist'),
    [('test-dummy', mock.DummyExpFamily(0.5, 0.5)),
     ('test-dirichlet',
      paddle.distribution.Dirichlet(paddle.to_tensor(parameterize.xrand()))),
     ('test-beta',
      paddle.distribution.Beta(paddle.to_tensor(parameterize.xrand()),
                               paddle.to_tensor(parameterize.xrand())))])
class TestExponentialFamilyException(unittest.TestCase):

    def test_entropy_exception(self):
        with self.assertRaises(NotImplementedError):
            paddle.distribution.ExponentialFamily.entropy(self.dist)


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