# 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 *
import pytest

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 xfail_when_nonstandard_decimal_separator


@xfail_when_nonstandard_decimal_separator
def test_group_adagrad():
    mx.random.seed(0)
    opt1 = mx.optimizer.contrib.GroupAdaGrad
    opt2 = mx.optimizer.contrib.GroupAdaGrad
    shapes = [(3, 4), [5, 6]]
    eps_options = [{}, {'epsilon': 1e-8}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    agg_options = [{}, {'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float32]:
        for options in itertools.product(eps_options, cg_options, rg_options, agg_options):
            kwarg = dict(wd=0.0)
            for option in options:
                kwarg.update(option)
            compare_optimizer(
                opt1(use_fused_step=False, **kwarg),
                opt2(use_fused_step=True, **kwarg),
                shapes,
                dtype)
            compare_optimizer(
                opt1(use_fused_step=False, **kwarg),
                opt2(use_fused_step=True, **kwarg),
                shapes,
                dtype,
                w_stype='row_sparse',
                g_stype='row_sparse')
            compare_optimizer(
                opt1(use_fused_step=False, **kwarg),
                opt2(use_fused_step=True, **kwarg),
                shapes,
                dtype,
                g_stype='row_sparse')

def _fn_noimpl(*args, **kwargs):
    raise NotImplementedError()

class _AdamLikeTestHelper:
    fn_update = _fn_noimpl
    fn_multi_update = _fn_noimpl
    fn_mp_update = _fn_noimpl
    fn_multi_mp_update = _fn_noimpl
    @staticmethod
    def ref_impl(m, v, weight, grad_rescale, beta1, beta2, lr, eta, wd, epsilon, clip_grad=-1):
        '''Returns (mean_ref, v_ref, weight_ref)'''
        raise NotImplementedError()
    @classmethod
    def run_test(cls, num_elem=1, aggregate=False):
        aggregate = aggregate or num_elem > 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(num_elem):
            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:
                cls.fn_multi_update(weight, grad, m, v,
                                     rescaled_grad, out=weight, **kwargs)
            else:
                cls.fn_update(weight[0], grad[0], m[0], v[0],
                               rescaled_grad, out=weight[0], **kwargs)
            # weights should remain unchanged
            for j in range(num_elem):
                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(num_elem):
            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:
                cls.fn_multi_mp_update(weight_fp16, grad_fp16, m, v, weight,
                                       rescaled_grad, out=weight_fp16, **kwargs)
            else:
                cls.fn_mp_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(num_elem):
                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(num_elem):
            grad_rescale.append(rescale_grad * grad[i])
            m_ref, v_ref, weight_ref = cls.ref_impl(
                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:
            cls.fn_multi_update(weight_test, grad, m, v,
                                rescale_grad, out=weight_test, **kwargs)
        else:
            cls.fn_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(num_elem):
            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(num_elem):
            grad_rescale.append(rescale_grad * grad_fp16[i].astype('float32'))
            m_ref, v_ref, weight_ref = cls.ref_impl(
                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:
            cls.fn_multi_mp_update(weight_fp16, grad_fp16, m, v, weight,
                                   rescale_grad, out=weight_fp16, **kwargs)
        else:
            cls.fn_mp_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(num_elem):
            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)

    def __call__(self):
        # Testing aggregated Adam update for one element
        self.run_test(1, aggregate=True)
        # Testing Adam update, if num_elem == 0, OR
        #         aggregated Adam update, if num_elem > 0
        for num_elem in reversed(range(6)):
            self.run_test(num_elem+1)

class _AdamWTestHelper(_AdamLikeTestHelper):
    fn_update = mx.nd.contrib.adamw_update
    fn_multi_update = mx.nd.contrib.multi_adamw_update
    fn_mp_update = mx.nd.contrib.mp_adamw_update
    fn_multi_mp_update = mx.nd.contrib.multi_mp_adamw_update
    @staticmethod
    def ref_impl(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

class _AdaBeliefTestHelper(_AdamLikeTestHelper):
    fn_update = mx.nd.contrib.adabelief_update
    fn_multi_update = mx.nd.contrib.multi_adabelief_update
    fn_mp_update = mx.nd.contrib.mp_adabelief_update
    fn_multi_mp_update = mx.nd.contrib.multi_mp_adabelief_update
    @staticmethod
    def ref_impl(m, v, weight, grad_rescale, beta1, beta2, lr, eta, wd, epsilon, clip_grad=-1):
        grad_rescale += wd * weight
        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-mean_ref)**2) + epsilon
        weight_ref = weight - eta * (lr * mean_ref / (v_ref.sqrt() + epsilon))
        return mean_ref, v_ref, weight_ref

@xfail_when_nonstandard_decimal_separator
@pytest.mark.serial
def test_adamw():
    _AdamWTestHelper()()

@xfail_when_nonstandard_decimal_separator
@pytest.mark.serial
def test_adabelief():
    _AdaBeliefTestHelper()()