# Copyright 2020 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""FTRL"""
from mindspore.ops import functional as F, composite as C, operations as P
from mindspore.common.parameter import Parameter
from mindspore.common import Tensor
import mindspore.common.dtype as mstype
from mindspore._checkparam import Validator as validator
from mindspore._checkparam import Rel
from .optimizer import Optimizer, apply_decay, grad_scale

ftrl_opt = C.MultitypeFuncGraph("ftrl_opt")


@ftrl_opt.register("Function", "Tensor", "Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor")
def _tensor_run_opt(opt, learning_rate, l1, l2, lr_power, linear, gradient, weight, moment):
    """Apply ftrl optimizer to the weight parameter."""
    success = True
    success = F.depend(success, opt(weight, moment, linear, gradient, learning_rate, l1, l2, lr_power))
    return success


def _check_param(initial_accum, learning_rate, lr_power, l1, l2, use_locking, loss_scale=1.0, weight_decay=0.0,
                 prim_name=None):
    """Check param."""
    validator.check_value_type("initial_accum", initial_accum, [float], prim_name)
    validator.check_number("initial_accum", initial_accum, 0.0, Rel.GE, prim_name)

    validator.check_value_type("learning_rate", learning_rate, [float], prim_name)
    validator.check_number("learning_rate", learning_rate, 0.0, Rel.GT, prim_name)

    validator.check_value_type("lr_power", lr_power, [float], prim_name)
    validator.check_number("lr_power", lr_power, 0.0, Rel.LE, prim_name)

    validator.check_value_type("l1", l1, [float], prim_name)
    validator.check_number("l1", l1, 0.0, Rel.GE, prim_name)

    validator.check_value_type("l2", l2, [float], prim_name)
    validator.check_number("l2", l2, 0.0, Rel.GE, prim_name)

    validator.check_value_type("use_locking", use_locking, [bool], prim_name)

    validator.check_value_type("loss_scale", loss_scale, [float], prim_name)
    validator.check_number("loss_scale", loss_scale, 1.0, Rel.GE, prim_name)

    validator.check_value_type("weight_decay", weight_decay, [float], prim_name)
    validator.check_number("weight_decay", weight_decay, 0.0, Rel.GE, prim_name)


class FTRL(Optimizer):
    """
    Implement the FTRL algorithm with ApplyFtrl Operator.

    FTRL is an online convex optimization algorithm that adaptively chooses its regularization function
    based on the loss functions. Refer to paper `Adaptive Bound Optimization for Online Convex Optimization
    <https://arxiv.org/abs/1002.4908>`_. Refer to paper `Ad Click Prediction: a View from the Trenches
    <https://www.eecs.tufts.edu/~dsculley/papers/ad-click-prediction.pdf>`_ for engineering document.

    Args:
        params (list[Parameter]): A list of parameter, which will be updated. The element in `params`
            should be Parameter.
        initial_accum (float): The starting value for accumulators, must be zero or positive values. Default: 0.1.
        learning_rate (float): The learning rate value, should be positive. Default: 0.001.
        lr_power (float): Learning rate power controls how the learning rate decreases during training, must be less
            than or equal to zero. Use fixed learning rate if lr_power is zero. Default: -0.5.
        l1 (float): l1 regularization strength, must be greater than or equal to zero. Default: 0.0.
        l2 (float): l2 regularization strength, must be greater than or equal to zero. Default: 0.0.
        use_locking (bool): If True use locks for update operation. Default: False.
        loss_scale (float): Value for the loss scale. It should be equal to or greater than 1.0. Default: 1.0.
        wegith_decay (float): Weight decay value to multiply weight, must be zero or positive value. Default: 0.0.

    Inputs:
        - **grads** (tuple[Tensor]) - The gradients of `params` in optimizer, the shape is as same as the `params`
          in optimizer.

    Outputs:
        tuple[Parameter], the updated parameters, the shape is the same as `params`.

    Examples:
        >>> net = Net()
        >>> loss = nn.SoftmaxCrossEntropyWithLogits()
        >>> opt = nn.FTRL(net.trainable_params())
        >>> model = Model(net, loss_fn=loss, optimizer=opt, metrics=None)
    """
    def __init__(self, params, initial_accum=0.1, learning_rate=0.001, lr_power=-0.5, l1=0.0, l2=0.0,
                 use_locking=False, loss_scale=1.0, weight_decay=0.0):
        super(FTRL, self).__init__(learning_rate, params)

        _check_param(initial_accum, learning_rate, lr_power, l1, l2, use_locking, loss_scale, weight_decay,
                     self.cls_name)
        self.moments = self.parameters.clone(prefix="moments", init=initial_accum)
        self.linear = self.parameters.clone(prefix="linear", init='zeros')
        self.l1 = l1
        self.l2 = l2
        self.lr_power = lr_power
        self.reciprocal_scale = 1.0 / loss_scale
        self.weight_decay = weight_decay
        self.decay_tf = tuple((lambda: True)() for x in self.parameters)
        self.hyper_map = C.HyperMap()
        self.opt = P.ApplyFtrl(use_locking=use_locking)
        self.one = Tensor(1, mstype.int32)

    def construct(self, grads):
        params = self.parameters
        moments = self.moments
        linear = self.linear
        if self.weight_decay > 0.0:
            grads = self.hyper_map(F.partial(apply_decay, self.weight_decay), self.decay_tf, params, grads)
        if self.reciprocal_scale != 1.0:
            grads = self.hyper_map(F.partial(grad_scale, self.reciprocal_scale), grads)
        lr = self.learning_rate
        success = self.hyper_map(F.partial(ftrl_opt, self.opt, lr, self.l1, self.l2, self.lr_power),
                                 linear, grads, params, moments)
        return success

[((24, 11, 24, 43), 'mindspore.ops.composite.MultitypeFuncGraph', 'C.MultitypeFuncGraph', ({(24, 32, 24, 42): '"""ftrl_opt"""'}, {}), "('ftrl_opt')", True, 'from mindspore.ops import functional as F, composite as C, operations as P\n'), ((38, 4, 38, 82), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(38, 31, 38, 46): '"""initial_accum"""', (38, 48, 38, 61): 'initial_accum', (38, 63, 38, 70): '[float]', (38, 72, 38, 81): 'prim_name'}, {}), "('initial_accum', initial_accum, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((39, 4, 39, 82), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(39, 27, 39, 42): '"""initial_accum"""', (39, 44, 39, 57): 'initial_accum', (39, 59, 39, 62): '(0.0)', (39, 64, 39, 70): 'Rel.GE', (39, 72, 39, 81): 'prim_name'}, {}), "('initial_accum', initial_accum, 0.0, Rel.GE, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((41, 4, 41, 82), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(41, 31, 41, 46): '"""learning_rate"""', (41, 48, 41, 61): 'learning_rate', (41, 63, 41, 70): '[float]', (41, 72, 41, 81): 'prim_name'}, {}), "('learning_rate', learning_rate, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((42, 4, 42, 82), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(42, 27, 42, 42): '"""learning_rate"""', (42, 44, 42, 57): 'learning_rate', (42, 59, 42, 62): '(0.0)', (42, 64, 42, 70): 'Rel.GT', (42, 72, 42, 81): 'prim_name'}, {}), "('learning_rate', learning_rate, 0.0, Rel.GT, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((44, 4, 44, 72), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(44, 31, 44, 41): '"""lr_power"""', (44, 43, 44, 51): 'lr_power', (44, 53, 44, 60): '[float]', (44, 62, 44, 71): 'prim_name'}, {}), "('lr_power', lr_power, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((45, 4, 45, 72), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(45, 27, 45, 37): '"""lr_power"""', (45, 39, 45, 47): 'lr_power', (45, 49, 45, 52): '(0.0)', (45, 54, 45, 60): 'Rel.LE', (45, 62, 45, 71): 'prim_name'}, {}), "('lr_power', lr_power, 0.0, Rel.LE, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((47, 4, 47, 60), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(47, 31, 47, 35): '"""l1"""', (47, 37, 47, 39): 'l1', (47, 41, 47, 48): '[float]', (47, 50, 47, 59): 'prim_name'}, {}), "('l1', l1, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((48, 4, 48, 60), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(48, 27, 48, 31): '"""l1"""', (48, 33, 48, 35): 'l1', (48, 37, 48, 40): '(0.0)', (48, 42, 48, 48): 'Rel.GE', (48, 50, 48, 59): 'prim_name'}, {}), "('l1', l1, 0.0, Rel.GE, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((50, 4, 50, 60), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(50, 31, 50, 35): '"""l2"""', (50, 37, 50, 39): 'l2', (50, 41, 50, 48): '[float]', (50, 50, 50, 59): 'prim_name'}, {}), "('l2', l2, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((51, 4, 51, 60), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(51, 27, 51, 31): '"""l2"""', (51, 33, 51, 35): 'l2', (51, 37, 51, 40): '(0.0)', (51, 42, 51, 48): 'Rel.GE', (51, 50, 51, 59): 'prim_name'}, {}), "('l2', l2, 0.0, Rel.GE, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((53, 4, 53, 77), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(53, 31, 53, 44): '"""use_locking"""', (53, 46, 53, 57): 'use_locking', (53, 59, 53, 65): '[bool]', (53, 67, 53, 76): 'prim_name'}, {}), "('use_locking', use_locking, [bool], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((55, 4, 55, 76), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(55, 31, 55, 43): '"""loss_scale"""', (55, 45, 55, 55): 'loss_scale', (55, 57, 55, 64): '[float]', (55, 66, 55, 75): 'prim_name'}, {}), "('loss_scale', loss_scale, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((56, 4, 56, 76), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(56, 27, 56, 39): '"""loss_scale"""', (56, 41, 56, 51): 'loss_scale', (56, 53, 56, 56): '(1.0)', (56, 58, 56, 64): 'Rel.GE', (56, 66, 56, 75): 'prim_name'}, {}), "('loss_scale', loss_scale, 1.0, Rel.GE, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((58, 4, 58, 80), 'mindspore._checkparam.Validator.check_value_type', 'validator.check_value_type', ({(58, 31, 58, 45): '"""weight_decay"""', (58, 47, 58, 59): 'weight_decay', (58, 61, 58, 68): '[float]', (58, 70, 58, 79): 'prim_name'}, {}), "('weight_decay', weight_decay, [float], prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((59, 4, 59, 80), 'mindspore._checkparam.Validator.check_number', 'validator.check_number', ({(59, 27, 59, 41): '"""weight_decay"""', (59, 43, 59, 55): 'weight_decay', (59, 57, 59, 60): '(0.0)', (59, 62, 59, 68): 'Rel.GE', (59, 70, 59, 79): 'prim_name'}, {}), "('weight_decay', weight_decay, 0.0, Rel.GE, prim_name)", True, 'from mindspore._checkparam import Validator as validator\n'), ((111, 25, 111, 37), 'mindspore.ops.composite.HyperMap', 'C.HyperMap', ({}, {}), '()', True, 'from mindspore.ops import functional as F, composite as C, operations as P\n'), ((112, 19, 112, 55), 'mindspore.ops.operations.ApplyFtrl', 'P.ApplyFtrl', (), '', True, 'from mindspore.ops import functional as F, composite as C, operations as P\n'), ((113, 19, 113, 42), 'mindspore.common.Tensor', 'Tensor', ({(113, 26, 113, 27): '1', (113, 29, 113, 41): 'mstype.int32'}, {}), '(1, mstype.int32)', False, 'from mindspore.common import Tensor\n'), ((124, 33, 124, 99), 'mindspore.ops.functional.partial', 'F.partial', ({(124, 43, 124, 51): 'ftrl_opt', (124, 53, 124, 61): 'self.opt', (124, 63, 124, 65): 'lr', (124, 67, 124, 74): 'self.l1', (124, 76, 124, 83): 'self.l2', (124, 85, 124, 98): 'self.lr_power'}, {}), '(ftrl_opt, self.opt, lr, self.l1, self.l2, self.lr_power)', True, 'from mindspore.ops import functional as F, composite as C, operations as P\n'), ((120, 35, 120, 76), 'mindspore.ops.functional.partial', 'F.partial', ({(120, 45, 120, 56): 'apply_decay', (120, 58, 120, 75): 'self.weight_decay'}, {}), '(apply_decay, self.weight_decay)', True, 'from mindspore.ops import functional as F, composite as C, operations as P\n'), ((122, 35, 122, 79), 'mindspore.ops.functional.partial', 'F.partial', ({(122, 45, 122, 55): 'grad_scale', (122, 57, 122, 78): 'self.reciprocal_scale'}, {}), '(grad_scale, self.reciprocal_scale)', True, 'from mindspore.ops import functional as F, composite as C, operations as P\n')]