# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 itertools
import numpy as np
import mxnet as mx
from mxnet.test_utils import *

curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
sys.path.insert(0, os.path.join(curr_path, '../unittest'))
from common import with_seed


# * GroupAdaGrad
class PyGroupAdaGrad(mx.optimizer.Optimizer):
    """The python reference of Group AdaGrad optimizer.

    Parameters
    ----------
    eps: float, optional
        Small value to avoid division by 0.

    """

    def __init__(self, eps=1e-5, **kwargs):
        super(PyGroupAdaGrad, self).__init__(**kwargs)
        self.float_stable_eps = eps

    def create_state(self, index, weight):
        assert len(weight.shape) == 2
        history = mx.nd.zeros(
            (weight.shape[0], 1), weight.context, stype=weight.stype)
        return history

    def update(self, index, weight, grad, state):
        self._update_count(index)
        lr = self._get_lr(index)
        wd = self._get_wd(index)
        assert wd == 0

        history = state
        grad = grad * self.rescale_grad
        if self.clip_gradient is not None:
            grad = mx.nd.clip(grad, -self.clip_gradient, self.clip_gradient)
        history[:] += mx.nd.mean(mx.nd.square(grad), axis=1, keepdims=True)
        div = lr * grad / mx.nd.sqrt(history + self.float_stable_eps)
        weight[:] -= div


def test_group_adagrad():
    mx.random.seed(0)
    opt1 = PyGroupAdaGrad
    opt2 = mx.optimizer.contrib.GroupAdaGrad
    shape = (3, 4)
    eps_options = [{}, {'eps': 1e-8}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    for dtype in [np.float32]:
        for options in itertools.product(eps_options, cg_options, rg_options):
            kwarg = dict(wd=0.0)
            for option in options:
                kwarg.update(option)
            compare_optimizer(
                opt1(**kwarg),
                opt2(**kwarg),
                shape,
                dtype,
                compare_states=False)
            compare_optimizer(
                opt1(**kwarg),
                opt2(**kwarg),
                shape,
                dtype,
                w_stype='row_sparse',
                g_stype='row_sparse',
                compare_states=False)
            compare_optimizer(
                opt1(**kwarg),
                opt2(**kwarg),
                shape,
                dtype,
                g_stype='row_sparse',
                compare_states=False)


@with_seed()
def test_adamw():
    def get_refs(m, v, weight, grad_rescale, beta1, beta2, lr, eta, wd, epsilon, clip_grad=-1):
        if clip_grad >= 0:
            grad_rescale = mx.nd.clip(grad_rescale, -clip_grad, clip_grad)

        mean_ref = beta1*m + (1-beta1)*grad_rescale
        v_ref = beta2*v + (1-beta2)*(grad_rescale**2)
        weight_ref = weight - eta * (lr * mean_ref / (v_ref.sqrt() + epsilon) + weight * wd)
        return mean_ref, v_ref, weight_ref

    def run_adamw_test(nElem=1, aggregate=False):
        aggregate = aggregate or nElem > 1
        rescale_factor = 10
        eta, lr, wd, epsilon = 1, 1, 0.1, 1e-8
        beta1, beta2 = 0.9, 0.999
        clip_gradient = np.random.uniform(rescale_factor, rescale_factor)
        weight, grad, m, v, etas, lrs, wds, weight_ref = [], [], [], [], [], [], [], []
        for i in range(nElem):
            shape = (np.random.randint(3, high=10), np.random.randint(3, high=10))
            weight.append(mx.nd.random.uniform(shape=shape))
            grad.append(mx.nd.random.uniform(-1.0, 1.0, shape=shape))
            m.append(mx.nd.random.uniform(shape=shape))
            v.append(mx.nd.random.uniform(shape=shape))
            etas.append(eta - 1 / np.random.uniform(9, 10))
            lrs.append(lr - 1 / np.random.uniform(9, 10))
            wds.append(wd - 1 / np.random.uniform(95, 105))
            weight_ref.append(weight[i].copy())

        if aggregate:
            kwargs = {'etas': etas, 'lrs': lrs, 'wds': wds}
        else:
            kwargs = {'eta': etas[0], 'lr': lrs[0], 'wd': wds[0]}

        kwargs.update([('epsilon', epsilon), ('beta1', beta1), ('beta2', beta2), ('clip_gradient', clip_gradient)])

        # Test 1: Update is skipped for rescale = nan scalar
        rescale_grad = mx.nd.array([rescale_factor])
        tested_grad = [rescale_grad * 0, rescale_grad * np.nan, rescale_grad * np.inf]
        tested_rescaled_grad = [np.nan]
        tested_rescaled_grad.extend(tested_grad)

        for rescaled_grad in tested_rescaled_grad:
            if aggregate:
                mx.nd.contrib.multi_adamw_update(weight, grad, m, v,
                                                 rescaled_grad, out=weight, **kwargs)
            else:
                mx.nd.contrib.adamw_update(weight[0], grad[0], m[0], v[0],
                                           rescaled_grad, out=weight[0], **kwargs)

            # weights should remain unchanged
            for j in range(nElem):
                assert_almost_equal(weight_ref[j], weight[j])


        # Test 2: Same as Test 1 for multi-precision update
        weight_fp16, grad_fp16, weight_fp16_refs = [], [], []
        for i in range(nElem):
            weight_fp16.append(weight[i].astype('float16'))
            grad_fp16.append(grad[i].astype('float16'))
            weight_fp16_refs.append(weight_fp16[i].copy())

        for rescaled_grad in tested_grad:
            if aggregate:
                mx.nd.contrib.multi_mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
                                                    rescaled_grad, out=weight_fp16, **kwargs)
            else:
                mx.nd.contrib.mp_adamw_update(weight_fp16[0], grad_fp16[0], m[0], v[0], weight[0],
                                              rescaled_grad, out=weight_fp16[0], **kwargs)

            # weights should remain unchanged
            for i in range(nElem):
                assert_almost_equal(weight_ref[i], weight[i])
                assert_almost_equal(weight_fp16_refs[i], weight_fp16[i])


        # Test 3: Reference normal update
        grad_rescale, weight_test, m_refs, v_refs, weight_refs = [], [], [], [], []
        for i in range(nElem):
            grad_rescale.append(rescale_grad * grad[i])
            m_ref, v_ref, weight_ref = get_refs(m[i], v[i], weight[i], grad_rescale[i], beta1, beta2, lrs[i], etas[i], wds[i], epsilon, clip_gradient)
            m_refs.append(m_ref)
            v_refs.append(v_ref)
            weight_refs.append(weight_ref)
            weight_test.append(weight[i].copy())

        # op normal update
        if aggregate:
            mx.nd.contrib.multi_adamw_update(weight_test, grad, m, v,
                                             rescale_grad, out=weight_test, **kwargs)
        else:
            mx.nd.contrib.adamw_update(weight_test[0], grad[0], m[0], v[0],
                                       rescale_grad, out=weight_test[0], **kwargs)

        # Compare results
        atol = 1e-4 if aggregate else 1e-5
        rtol = 1e-4 if aggregate else None
        for i in range(nElem):
            assert_almost_equal(weight_refs[i], weight_test[i], rtol=rtol, atol=atol)
            assert_almost_equal(m_refs[i], m[i], rtol=rtol, atol=atol)
            assert_almost_equal(v_refs[i], v[i], atol=atol)


        # Test 4: Reference normal multi-precision update
        grad_rescale, m_refs, v_refs, weight_refs, weight_fp16_refs = [], [], [], [], []
        for i in range(nElem):
            grad_rescale.append(rescale_grad * grad_fp16[i].astype('float32'))
            m_ref, v_ref, weight_ref = get_refs(m[i], v[i], weight[i], grad_rescale[i], beta1, beta2, lrs[i], etas[i], wds[i], epsilon, clip_gradient)
            m_refs.append(m_ref)
            v_refs.append(v_ref)
            weight_refs.append(weight_ref)
            weight_fp16_refs.append(weight_ref.astype('float16'))

        # op normal multi-precision update
        if aggregate:
            mx.nd.contrib.multi_mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
                                                rescale_grad, out=weight_fp16, **kwargs)
        else:
            mx.nd.contrib.mp_adamw_update(weight_fp16[0], grad_fp16[0], m[0], v[0], weight[0],
                                          rescale_grad, out=weight_fp16[0], **kwargs)

        # Compare results
        for i in range(nElem):
            assert_almost_equal(m_refs[i], m[i], rtol=rtol, atol=atol)
            assert_almost_equal(v_refs[i], v[i], atol=atol)
            assert_almost_equal(weight_refs[i], weight[i], rtol=rtol, atol=atol)
            assert_almost_equal(weight_fp16_refs[i], weight_fp16[i], rtol=1e-3, atol=atol)

    # Testing aggregated Adam update for one element
    run_adamw_test(1, aggregate=True)

    # Testing Adam update, if nElem = 0, OR
    #         aggregated Adam update, if nElem > 0
    for nElem in range(6):
        run_adamw_test(nElem+1)

if __name__ == '__main__':
    import nose
    nose.runmodule()