# Copyright (c) 2018 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 unittest
import logging
import six


class TestBase(unittest.TestCase):
    def main(self,
             network_func,
             iter=10,
             iter_per_pe=10,
             use_gpu=True,
             use_experimental_executor=False):
        if use_gpu and not fluid.core.is_compiled_with_cuda():
            logging.warning(
                "Paddle is not compiled with CUDA, skip GPU unittests")
            return

        main_prog = fluid.Program()
        startup_prog = fluid.Program()
        scope = fluid.Scope()
        with fluid.program_guard(main_prog, startup_prog):
            with fluid.scope_guard(scope):
                loss = network_func()
                fluid.Executor(
                    fluid.CUDAPlace(0)
                    if use_gpu else fluid.CPUPlace()).run(startup_prog)

        for _ in six.moves.xrange(iter):
            exe_strategy = fluid.ExecutionStrategy()
            exe_strategy._dry_run = True
            exe_strategy.use_experimental_executor = use_experimental_executor
            pe = fluid.ParallelExecutor(
                use_cuda=use_gpu,
                loss_name=loss.name,
                main_program=main_prog,
                exec_strategy=exe_strategy)
            for _ in six.moves.xrange(iter_per_pe):
                pe.run([])


class TestMNISTDryRun(TestBase):
    def test_mnist_dry_run(self):
        for use_gpu in (False, True):
            for use_experimental_executor in (False, True):
                self.main(
                    network_func=TestMNISTDryRun.network_func,
                    use_gpu=use_gpu,
                    use_experimental_executor=use_experimental_executor)

    @staticmethod
    def network_func():
        img = fluid.layers.data(name='img', shape=[784], dtype='float32')
        label = fluid.layers.data(name='label', shape=[1], dtype='int64')
        hidden = img
        for _ in six.moves.xrange(10):
            hidden = fluid.layers.fc(input=img, size=200, act='tanh')
        prediction = fluid.layers.fc(input=hidden, size=10, act='softmax')
        loss = fluid.layers.cross_entropy(input=prediction, label=label)
        avg_loss = fluid.layers.mean(loss)
        fluid.optimizer.Adam().minimize(avg_loss)
        return avg_loss


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