"""Weight updating functions"""
import math
import pickle
import logging
from .ndarray import NDArray, zeros, clip, sqrt
from .ndarray import sgd_update, sgd_mom_update, adam_update, rmsprop_update, rmspropalex_update
from .random import normal


class Optimizer(object):
    """The base class inherited by all optimizers.

    Parameters
    ----------
    rescale_grad : float, optional
        Multiply the gradient with ``rescale_grad`` before updating. Often
        choose to be ``1.0/batch_size``.

    param_idx2name : dict from int to string, optional
        A dictionary that maps int index to string name.

    clip_gradient : float, optional
        Clip the gradient by projecting onto the box ``[-clip_gradient, clip_gradient]``.

    learning_rate : float, optional
        The initial learning rate.

    lr_scheduler : LRScheduler, optional
        The learning rate scheduler.

    wd : float, optional
        The weight decay (or L2 regularization) coefficient. Modifies objective
        by adding a penalty for having large weights.

    sym: Symbol, optional
        The Symbol this optimizer is applying to.

    begin_num_update : int, optional
        The initial number of updates
    """
    def __init__(self, rescale_grad=1., param_idx2name=None, wd=0.,
                 clip_gradient=None, learning_rate=0.01,
                 lr_scheduler=None, sym=None, begin_num_update=0):
        self.rescale_grad = rescale_grad
        self.lr = learning_rate
        self.lr_scheduler = lr_scheduler
        if lr_scheduler is not None:
            self.lr_scheduler.base_lr = learning_rate

        self.wd = wd
        self.lr_mult = {}
        self.wd_mult = {}
        self.begin_num_update = begin_num_update
        self.num_update = begin_num_update
        self._index_update_count = {}
        self.clip_gradient = clip_gradient

        if param_idx2name is None:
            param_idx2name = {}
        assert isinstance(param_idx2name, dict), \
            'param_idx2name should be a dict of param indexes to names.'
        self.idx2name = param_idx2name.copy()
        self.sym = sym

        self.set_lr_mult({})
        self.set_wd_mult({})

    opt_registry = {}

    @staticmethod
    def register(klass):
        """Register a new optimizer.

        Once an optimizer is registered, we can create an instance of this
        optimizer with ``create_optimizer`` later.

        Examples
        --------

        >>> @mx.optimizer.Optimizer.register
        ... class MyOptimizer(mx.optimizer.Optimizer):
        ...     pass
        >>> optim = mx.optimizer.Optimizer.create_optimizer('MyOptimizer')
        >>> print(type(optim))
        <class '__main__.MyOptimizer'>
        """
        assert(isinstance(klass, type))
        name = klass.__name__.lower()
        if name in Optimizer.opt_registry:
            logging.warning('WARNING: New optimizer %s.%s is overriding '
                            'existing optimizer %s.%s',
                            klass.__module__, klass.__name__,
                            Optimizer.opt_registry[name].__module__,
                            Optimizer.opt_registry[name].__name__)
        Optimizer.opt_registry[name] = klass
        return klass

    @staticmethod
    def create_optimizer(name, **kwargs):
        """Instantiate an optimizer with a given name and kwargs.

        Notes
        -----
        We can use the alias ``create`` for ``Optimizer.create_optimizer``

        Parameters
        ----------
        name: str
            Name of the optimizer. Should be the name
            of a subclass of Optimizer. Case insensitive.

        kwargs: dict
            Parameters for the optimizer.

        Returns
        -------
        Optimizer
            An instantiated optimizer.

        Examples
        --------
        >>> sgd = mx.optimizer.Optimizer.create_optimizer('sgd')
        >>> type(sgd)
        <class 'mxnet.optimizer.SGD'>
        >>> adam = mx.optimizer.create('adam', learning_rate=.1)
        >>> type(adam)
        <class 'mxnet.optimizer.Adam'>
        """
        if name.lower() in Optimizer.opt_registry:
            return Optimizer.opt_registry[name.lower()](**kwargs)
        else:
            raise ValueError('Cannot find optimizer %s' % name)


    def create_state(self, index, weight):
        """Create auxiliary state for a given weight

        Some optimizers require additional states, e.g. as momentum, in addition
        to gradients in order to update weights. This function creates state
        for a given weight which will be used in ``update``. This function is
        called only once for each weight.

        Parameters
        ----------
        index : int
            An unique index to identify the weight.
        weight : NDArray
            The weight

        Returns
        -------
        state : any obj
            The state associated with the weight.
        """

    def update(self, index, weight, grad, state):
        """Update the weight given the corresponding gradient and state.

        Parameters
        ----------
        index : int
            An unique index to identify the weight.
        weight : NDArray
            The weight
        grad : NDArray
            The gradient of the objective with respect to this weight.
        state : any obj
            The state associated with this weight.
        """
        raise NotImplementedError()

    def set_lr_scale(self, args_lrscale): # pylint: disable=unused-argument
        """[DEPRECATED] set lr scale. Use set_lr_mult instead."""
        raise DeprecationWarning

    def set_lr_mult(self, args_lr_mult):
        """Set individual learning rate for each weight.

        Parameters
        ----------
        args_lr_mult : dict of string/int to float
            Set the lr multipler for name/index to float.
            Setting multipler by index is supported for backward compatibility,
            but we recommend using name and symbol.
        """
        self.lr_mult = {}
        if self.sym is not None:
            attr = self.sym.attr_dict()
            for name in self.sym.list_arguments():
                if name in attr and '__lr_mult__' in attr[name]:
                    self.lr_mult[name] = float(attr[name]['__lr_mult__'])
        self.lr_mult.update(args_lr_mult)

    def set_wd_mult(self, args_wd_mult):
        """Set individual weight decay for each weight.

        By default wd multipler is 0 for all params whose name doesn't
        end with _weight, if param_idx2name is provided.

        Parameters
        ----------
        args_wd_mult : dict of string/int to float
            Set the wd multipler for name/index to float.
            Setting multipler by index is supported for backward compatibility,
            but we recommend using name and symbol.
        """
        self.wd_mult = {}
        for n in self.idx2name.values():
            if not (n.endswith('_weight') or n.endswith('_gamma')):
                self.wd_mult[n] = 0.0
        if self.sym is not None:
            attr = self.sym.attr_dict()
            for name in self.sym.list_arguments():
                if name in attr and '__wd_mult__' in attr[name]:
                    self.wd_mult[name] = float(attr[name]['__wd_mult__'])
        self.wd_mult.update(args_wd_mult)

    def _update_count(self, index):
        """Update num_update

        Parameters:
        index : int
            The index to be updated.
        """
        if index not in self._index_update_count:
            self._index_update_count[index] = self.begin_num_update
        self._index_update_count[index] += 1
        self.num_update = max(self._index_update_count[index], self.num_update)

    def _get_lr(self, index):
        """Get the learning rate given the index of the weight.

        Parameters
        ----------
        index : int
            The index corresponding to the weight.

        Returns
        -------
        lr : float
            Learning rate for this index.
        """
        if self.lr_scheduler is not None:
            lr = self.lr_scheduler(self.num_update)
        else:
            lr = self.lr

        if index in self.lr_mult:
            lr *= self.lr_mult[index]
        elif index in self.idx2name:
            lr *= self.lr_mult.get(self.idx2name[index], 1.0)
        return lr

    def _get_wd(self, index):
        """get weight decay for index.
        Returns 0 for non-weights if the name of weights are provided for __init__.

        Parameters
        ----------
        index : int
            The index for weight.

        Returns
        -------
        wd : float
            Weight decay for this index.
        """
        wd = self.wd
        if index in self.wd_mult:
            wd *= self.wd_mult[index]
        elif index in self.idx2name:
            wd *= self.wd_mult.get(self.idx2name[index], 1.0)
        return wd

# convenience wrapper for Optimizer.Register
register = Optimizer.register   # pylint: disable=invalid-name

@register
class SGD(Optimizer):
    """The SGD optimizer with momentum and weight decay.

    The optimizer updates the weight by:

      state = momentum * state + lr * rescale_grad * clip(grad, clip_gradient) + wd * weight
      weight = weight - state

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.Optimizer`:

    Parameters
    ----------
    momentum : float, optional
       The momentum value.
    """
    def __init__(self, momentum=0.0, **kwargs):
        super(SGD, self).__init__(**kwargs)
        self.momentum = momentum
        self.kwargs = {'rescale_grad': self.rescale_grad}
        if self.momentum > 0:
            self.kwargs['momentum'] = self.momentum
        if self.clip_gradient:
            self.kwargs['clip_gradient'] = self.clip_gradient

    def create_state(self, index, weight):
        if self.momentum == 0.0:
            return None
        else:
            return zeros(weight.shape, weight.context, dtype=weight.dtype)

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

        if state:
            sgd_mom_update(weight, grad, state, out=weight,
                           lr=lr, wd=wd, **self.kwargs)
        else:
            sgd_update(weight, grad, out=weight,
                       lr=lr, wd=wd, **self.kwargs)

@register
class DCASGD(Optimizer):
    """The DCASGD optimizer

    This class implements the optimizer described in *Asynchronous Stochastic Gradient Descent with
    Delay Compensation for Distributed Deep Learning*, available at https://arxiv.org/abs/1609.08326

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.Optimizer`:

    Parameters
    ----------
    momentum : float, optional
       The momentum value.

    lamda : float, optional
       Scale DC value.
    """
    def __init__(self, momentum=0.0, lamda=0.04, **kwargs):
        super(DCASGD, self).__init__(**kwargs)
        self.momentum = momentum
        self.weight_previous = {}
        self.lamda = lamda

    def create_state(self, index, weight):
        if self.momentum == 0.0:
            return (None,
                    weight.copy())  # previous weight
        else:
            return (zeros(weight.shape, weight.context, dtype=weight.dtype), # momentum
                    weight.copy())  # previous weight

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

        grad = grad * self.rescale_grad
        if self.clip_gradient is not None:
            grad = clip(grad, -self.clip_gradient, self.clip_gradient)

        mom, previous_weight = state
        if mom:
            mom[:] *= self.momentum
            mom[:] += -lr * (grad + wd * weight + self.lamda \
                      * grad * grad * (weight - previous_weight))
        else:
            assert(self.momentum == 0.0)
            mom = -lr * (grad + wd * weight + self.lamda \
                      * grad * grad * (weight - previous_weight))
        previous_weight[:] = weight
        weight[:] += mom

@register
class NAG(SGD):
    """Nesterov accelerated SGD.

    This optimizer updates each weight by:

        state = momentum * state + grad + wd * weight
        weight = weight - (lr * (grad + momentum * state))

    This optimizer accepts the same arguments as :class:`.SGD`.
    """
    def __init__(self, **kwargs):
        super(NAG, self).__init__(**kwargs)

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

        grad = grad * self.rescale_grad
        if self.clip_gradient is not None:
            grad = clip(grad, -self.clip_gradient, self.clip_gradient)

        if state:
            mom = state
            mom[:] *= self.momentum
            grad += wd * weight
            mom[:] += grad
            grad[:] += self.momentum * mom
            weight[:] += -lr * grad
        else:
            assert self.momentum == 0.0
            weight[:] += -lr * (grad + wd * weight)

@register
class SGLD(Optimizer):
    """Stochastic Gradient Riemannian Langevin Dynamics.

    This class implements the optimizer described in the paper *Stochastic Gradient
    Riemannian Langevin Dynamics on the Probability Simplex*, available at
    https://papers.nips.cc/paper/4883-stochastic-gradient-riemannian-langevin-dynamics-on-the-probability-simplex.pdf

    """
    def __init__(self, **kwargs):
        super(SGLD, self).__init__(**kwargs)

    def create_state(self, index, weight):
        return None

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

        grad = grad * self.rescale_grad
        if self.clip_gradient is not None:
            grad = clip(grad, -self.clip_gradient, self.clip_gradient)
        weight[:] += - lr/2 * (grad + wd * weight) + normal(0, math.sqrt(lr),
                                                            weight.shape, weight.context)


@register  # pylint: disable=invalid-name
class ccSGD(SGD):
    """[Deprecated] Same as sgd. Left here for backward compatibility."""
    def __init__(self, *args, **kwargs):
        super(ccSGD, self).__init__(*args, **kwargs)

@register
class Adam(Optimizer):
    """The Adam optimizer.

    This class implements the optimizer described in *Adam: A Method for
    Stochastic Optimization*, available at http://arxiv.org/abs/1412.6980

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.Optimizer`:

    Parameters
    ----------
    beta1 : float, optional
        Exponential decay rate for the first moment estimates.
    beta2 : float, optional
        Exponential decay rate for the second moment estimates.
    epsilon : float, optional
        Small value to avoid divided by 0.
    """
    def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
                 **kwargs):
        super(Adam, self).__init__(learning_rate=learning_rate, **kwargs)
        self.beta1 = beta1
        self.beta2 = beta2
        self.kwargs = {'beta1': beta1, 'beta2': beta2, 'epsilon': epsilon,
                       'rescale_grad': self.rescale_grad}
        if self.clip_gradient:
            self.kwargs['clip_gradient'] = self.clip_gradient

    def create_state(self, index, weight):
        return (zeros(weight.shape, weight.context, dtype=weight.dtype),  # mean
                zeros(weight.shape, weight.context, dtype=weight.dtype))  # variance

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

        t = self._index_update_count[index]
        coef1 = 1. - self.beta1**t
        coef2 = 1. - self.beta2**t
        lr *= math.sqrt(coef2)/coef1
        mean, var = state
        adam_update(weight, grad, mean, var, out=weight,
                    lr=lr, wd=wd, **self.kwargs)

@register
class AdaGrad(Optimizer):
    """AdaGrad optimizer

    This calss implements the AdaGrad optiizer described in *Adaptive Subgradient
    Methods for Online Learning and Stochastic Optimization*, and available at
    http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.Optimizer`:

    Parameters
    ----------
    eps: float, optional
        Small value to avoid division by 0.
    """
    def __init__(self, eps=1e-7, **kwargs):
        super(AdaGrad, self).__init__(**kwargs)
        self.float_stable_eps = eps

    def create_state(self, index, weight):
        return zeros(weight.shape, weight.context)  # history

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

        grad = grad * self.rescale_grad
        if self.clip_gradient is not None:
            grad = clip(grad, -self.clip_gradient, self.clip_gradient)
        history = state
        history[:] += (grad * grad)
        weight[:] += -lr * (grad / sqrt(history + self.float_stable_eps) + wd * weight)

@register
class RMSProp(Optimizer):
    """The RMSProp optimizer.

    Two versions of RMSProp are implemented:

    If ``centered=False``, we follow
    http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf by
    Tieleman & Hinton, 2012.

    If ``centered=True``, we follow http://arxiv.org/pdf/1308.0850v5.pdf (38)-(45)
    by Alex Graves, 2013.

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.Optimizer`:

    Parameters
    ----------
    gamma1: float, optional
        Decay factor of moving average for ``gradient^2``.
    gamma2: float, optional
        A "momentum" factor. Only used if ``centered=True``.
    epsilon : float, optional
        Small value to avoid division by 0.
    centered : bool, optional
        Use Graves' or Tieleman & Hinton's version of RMSProp.
    clip_weights : float, optional
        clip weights into range ``[-clip_weights, clip_weights]``
    """
    def __init__(self, learning_rate=0.001, gamma1=0.9, gamma2=0.9,
                 epsilon=1e-8, centered=False, clip_weights=None, **kwargs):
        super(RMSProp, self).__init__(learning_rate=learning_rate, **kwargs)
        self.gamma1 = gamma1
        self.gamma2 = gamma2
        self.centered = centered
        self.clip_weights = clip_weights
        self.kwargs = {'gamma1': gamma1, 'epsilon': epsilon,
                       'rescale_grad': self.rescale_grad}
        if self.centered:
            self.kwargs['gamma2'] = gamma2
        if self.clip_gradient:
            self.kwargs['clip_gradient'] = self.clip_gradient
        if self.clip_weights:
            self.kwargs['clip_weights'] = self.clip_weights

    def create_state(self, index, weight):
        if self.centered:
            return (
                zeros(weight.shape, weight.context),  # n
                zeros(weight.shape, weight.context),  # g
                zeros(weight.shape, weight.context))  # delta
        else:
            return (zeros(weight.shape, weight.context), )  # n

    def update(self, index, weight, grad, state):
        assert(isinstance(weight, NDArray))
        assert(isinstance(grad, NDArray))
        lr = self._get_lr(index)
        wd = self._get_wd(index)
        self._update_count(index)
        if not self.centered:
            (n, ) = state
            rmsprop_update(
                weight, grad, n, out=weight, lr=lr, wd=wd, **self.kwargs)
        else:
            n, g, delta = state
            rmspropalex_update(weight, grad, n, g, delta, out=weight,
                               lr=lr, wd=wd, **self.kwargs)

@register
class AdaDelta(Optimizer):
    """The AdaDelta optimizer.

    This class implements AdaDelta, an optimizer described in  *ADADELTA: An adaptive
    learning rate method*, available at https://arxiv.org/abs/1212.5701

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.Optimizer`:

    Parameters
    ----------
    rho: float
        Decay rate for both squared gradients and delta.
    epsilon : float
        Small value to avoid division by 0.
    """
    def __init__(self, rho=0.90, epsilon=1e-5, **kwargs):
        super(AdaDelta, self).__init__(**kwargs)
        self.rho = rho
        self.epsilon = epsilon

    def create_state(self, index, weight):
        return (zeros(weight.shape, weight.context),  # accumulated g
                zeros(weight.shape, weight.context))  # accumulated delta

    def update(self, index, weight, grad, state):
        assert(isinstance(weight, NDArray))
        assert(isinstance(grad, NDArray))
        wd = self._get_wd(index)
        self._update_count(index)

        # preprocess grad
        grad *= self.rescale_grad
        if self.clip_gradient is not None:
            grad = clip(grad, -self.clip_gradient, self.clip_gradient)

        # accumulated g and delta initlization
        acc_g, acc_delta = state

        # update g, delta
        acc_g[:] = self.rho * acc_g + (1. - self.rho) * grad * grad
        current_delta = sqrt(acc_delta + self.epsilon) / sqrt(acc_g + self.epsilon) * grad
        acc_delta[:] = self.rho * acc_delta + (1. - self.rho) * current_delta * current_delta

        # update weight
        weight[:] -= current_delta + wd * weight

@register
class Test(Optimizer):
    def __init__(self, **kwargs):
        super(Test, self).__init__(**kwargs)

    def create_state(self, index, weight):
        """Create a state to duplicate weight"""
        return zeros(weight.shape, weight.context)

    def update(self, index, weight, grad, state):
        """performs w += rescale_grad * grad"""
        weight[:] += grad * self.rescale_grad
        state[:] = weight

# backward compatibility wrapper for Optimizer.CreateOptimizer
create = Optimizer.create_optimizer  # pylint: disable=invalid-name

class Updater(object):
    """Updater for kvstore."""
    def __init__(self, optimizer):
        self.optimizer = optimizer
        self.states = {}

    def __call__(self, index, grad, weight):
        """Update weight given gradient and index."""
        if index not in self.states:
            self.states[index] = self.optimizer.create_state(index, weight)
        self.optimizer.update(index, weight, grad, self.states[index])

    def set_states(self, states):
        """Set updater states."""
        self.states = pickle.loads(states)

    def get_states(self):
        """Get updater states."""
        return pickle.dumps(self.states)

def get_updater(optimizer):
    """Return a clossure of the updater needed for kvstore.

    Parameters
    ----------
    optimizer: Optimizer
         The optimizer.

    Returns
    -------
    updater: function
         The clossure of the updater.
    """
    return Updater(optimizer)