# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
from typing import Iterable, Tuple, Union

from ..core.tensor.tensor import Tensor
from ..tensor import Parameter, tensor
from .optimizer import Optimizer


class Adam(Optimizer):
    r"""
    Implements Adam algorithm proposed in `"Adam: A Method for Stochastic Optimization" <https://arxiv.org/abs/1412.6980>`_.

    :param params: iterable of parameters to optimize or dicts defining
            parameter groups.
    :param lr: learning rate.
    :param betas: coefficients used for computing running averages of gradient
        and its square. Default: (0.9, 0.999)
    :param eps: term added to the denominator to improve numerical stability
        Default: 1e-8
    :param weight_decay: weight decay (L2 penalty). Default: 0
    """

    def __init__(
        self,
        params: Union[Iterable[Parameter], dict],
        lr: float,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0.0,
    ):
        if lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if weight_decay < 0.0:
            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))

        defaults = dict(lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
        super().__init__(params, defaults)

    def _create_state(self, param_group):
        for param in param_group["params"]:
            self._add_state(param, "exp_avg")
            self._add_state(param, "exp_avg_sq")
            self._add_state(param, "step", initializer=0.0)

    def _updates(self, param_group):
        lr = param_group["lr"]
        weight_decay = param_group["weight_decay"]
        eps = param_group["eps"]
        beta0, beta1 = param_group["betas"]

        # since `conver_inputs` is disabled for param updates,
        # scalar should be explicitly tansforred to tensor
        _lr = tensor([lr])
        _weight_decay = tensor([weight_decay])
        _eps = tensor([eps])
        _beta0, _beta1 = tensor([beta0]), tensor([beta1])

        c1 = tensor([1.0])
        c05 = tensor([0.5])
        for param in param_group["params"]:

            if param.grad is None:
                continue

            grad = param.grad
            if weight_decay != 0.0:
                grad += param * _weight_decay

            states = self._state[param]
            step = states["step"]
            step += c1
            exp_avg = states["exp_avg"]
            exp_avg_sq = states["exp_avg_sq"]
            exp_avg = _beta0 * exp_avg + grad * (c1 - _beta0)
            exp_avg_sq = _beta1 * exp_avg_sq + (c1 - _beta1) * (grad * grad)

            delta = (exp_avg / (c1 - _beta0 ** step)) / (
                (exp_avg_sq / (c1 - _beta1 ** step)) ** c05 + _eps
            )
            param -= _lr * delta

            # not inplace change, need to update underlying tensor handler in state
            states["exp_avg"]._reset(exp_avg)
            states["exp_avg_sq"]._reset(exp_avg_sq)