# Copyright (c) 2022 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 sys

sys.path.append('..')
import unittest
import paddle
import numpy as np
from paddle.fluid.layer_helper import LayerHelper
from collections import OrderedDict


def run_momentum_op(params,
                    grads,
                    velocitys,
                    master_params,
                    learning_rate,
                    place,
                    multi_precision,
                    mu=0.9,
                    rescale_grad=0.01,
                    use_merged=False,
                    use_nesterov=True):
    assert len(params) == len(grads)
    assert len(params) == len(velocitys)
    if multi_precision:
        assert len(params) == len(master_params)
    op_type = 'merged_momentum' if use_merged else 'momentum'
    main = paddle.static.Program()
    startup = paddle.static.Program()
    with paddle.static.program_guard(main, startup):
        helper = LayerHelper(op_type, **locals())

        param_vars = [
            helper.create_variable(persistable=True,
                                   shape=p.shape,
                                   dtype=p.dtype) for p in params
        ]
        grad_vars = [
            helper.create_variable(shape=g.shape, dtype=g.dtype) for g in grads
        ]
        velocity_vars = [
            helper.create_variable(persistable=True,
                                   shape=v.shape,
                                   dtype=v.dtype) for v in velocitys
        ]
        lr_var = helper.create_variable(persistable=True,
                                        shape=learning_rate.shape,
                                        dtype=learning_rate.dtype)

        feed_dict = OrderedDict()

        feed_dict.update(
            OrderedDict([(p_var.name, p_val)
                         for p_var, p_val in zip(param_vars, params)]))
        feed_dict.update(
            OrderedDict([(v_var.name, v_val)
                         for v_var, v_val in zip(velocity_vars, velocitys)]))
        fetch_list = list(feed_dict.keys())

        feed_dict.update(
            OrderedDict([(g_var.name, g_val)
                         for g_var, g_val in zip(grad_vars, grads)]))
        feed_dict.update({lr_var.name: learning_rate})

        if multi_precision:
            master_param_vars = [
                helper.create_variable(persistable=True,
                                       shape=p.shape,
                                       dtype=p.dtype) for p in master_params
            ]
            feed_dict.update(
                OrderedDict([
                    (mp_var.name, mp_val)
                    for mp_var, mp_val in zip(master_param_vars, master_params)
                ]))
            # CPUPlace does not use MasterParam
            if isinstance(place, paddle.CUDAPlace):
                fetch_list = fetch_list + [
                    mp_var.name for mp_var in master_param_vars
                ]
        else:
            master_param_vars = None

        if not use_merged:
            for i, (p, g,
                    v) in enumerate(zip(param_vars, grad_vars, velocity_vars)):
                inputs = {
                    'Param': p,
                    'Grad': g,
                    'Velocity': v,
                    'LearningRate': lr_var,
                }
                outputs = {'ParamOut': p, 'VelocityOut': v}
                if multi_precision:
                    inputs['MasterParam'] = master_param_vars[i]
                    outputs['MasterParamOut'] = master_param_vars[i]
                attrs = {
                    'mu': mu,
                    'multi_precision': multi_precision,
                    'rescale_grad': rescale_grad,
                    'use_nesterov': use_nesterov,
                    'regularization_method': 'l2_decay',
                    'regularization_coeff': 2.0,
                }
                helper.append_op(type=op_type,
                                 inputs=inputs,
                                 outputs=outputs,
                                 attrs=attrs)
        else:
            inputs = {
                'Param': param_vars,
                'Grad': grad_vars,
                'Velocity': velocity_vars,
                'LearningRate': lr_var,
            }
            outputs = {'ParamOut': param_vars, 'VelocityOut': velocity_vars}
            if multi_precision:
                inputs['MasterParam'] = master_param_vars
                outputs['MasterParamOut'] = master_param_vars
            attrs = {
                'mu':
                mu,
                'multi_precision':
                multi_precision,
                'rescale_grad':
                rescale_grad,
                'use_nesterov':
                use_nesterov,
                'regularization_method':
                ['l2_decay' for i in range(len(param_vars))],
                'regularization_coeff': [2.0 for i in range(len(param_vars))],
            }
            helper.append_op(type=op_type,
                             inputs=inputs,
                             outputs=outputs,
                             attrs=attrs)

    exe = paddle.static.Executor(place)
    with paddle.static.scope_guard(paddle.static.Scope()):
        exe.run(startup)
        return exe.run(main, feed=feed_dict, fetch_list=fetch_list)


class TestMergedMomentumBase(unittest.TestCase):

    def setUp(self):
        paddle.enable_static()
        self.shapes = [[3, 4], [2, 7], [5, 6], [7, 8]]
        self.seed = 10
        self.place = paddle.fluid.XPUPlace(0)
        self.__class__.use_xpu = True

    def gen_rand_data(self, shapes, dtype):
        return [np.random.random(s).astype(dtype) for s in shapes]

    def prepare_data(self, shapes, multi_precision, seed, dtype, place):
        np.random.seed(seed)
        params = self.gen_rand_data(shapes, dtype)
        grads = self.gen_rand_data(shapes, dtype)
        velocitys = self.gen_rand_data(shapes, dtype)
        learning_rate = self.gen_rand_data([[1]], np.float32)[0]
        if multi_precision:
            master_params = [p.astype(dtype) for p in params]
        else:
            master_params = None
        return params, grads, velocitys, master_params, learning_rate

    def check_with_place(self, place, dtype, multi_precision=False):
        params, grads, velocitys, master_params, learning_rate = self.prepare_data(
            self.shapes, multi_precision, self.seed, dtype, place)

        def run_op(use_nesterov, use_merged):
            # NPU Momentum Op does not support rescale_grad
            rescale_grad = 1.0
            return run_momentum_op(params,
                                   grads,
                                   velocitys,
                                   master_params,
                                   learning_rate,
                                   place,
                                   multi_precision,
                                   rescale_grad=rescale_grad,
                                   use_merged=use_merged,
                                   use_nesterov=use_nesterov)

        outs1 = run_op(use_nesterov=True, use_merged=True)
        outs2 = run_op(use_nesterov=True, use_merged=False)
        self.assertEqual(len(outs1), len(outs2))
        for i, (out1, out2) in enumerate(zip(outs1, outs2)):
            np.testing.assert_allclose(out1, out2, atol=1e-7)
        outs3 = run_op(use_nesterov=False, use_merged=True)
        outs4 = run_op(use_nesterov=False, use_merged=False)
        self.assertEqual(len(outs3), len(outs4))
        for j, (out3, out4) in enumerate(zip(outs3, outs4)):
            np.testing.assert_allclose(out3, out4, atol=1e-7)


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