from __future__ import division

from chainer.training import extension


class WarmupShift(extension.Extension):

    """Trainer extension to change an optimizer attribute evenly at the
    begining of one training.

    For example, suppose that this extension is called at every iteration,
    and warmup_start = x , init = y, warmup_iter = t.
    Then this extension will set the corresponding attribute to from
    ``x`` to ``y`` evenly in first ``t`` iterations.

    This extension is also called before the training loop starts by default.

    Args:
        attr (str): Name of the optimizer attribute to adjust.
        warmup_start (float): the value of the attr at the begining
            of one training.
        init (float): the value of the attr after warm up iterations.
        warmup_iter (int): the number of the iterations in which the
            attr changes from ``warmup_start`` to ``init``.
        optimizer (~chainer.Optimizer): Target optimizer object.
            If it is None, the main optimizer of the trainer is used.

    """

    def __init__(self, attr, warmup_start, warmup_iter, init,
                 optimizer=None):
        self._attr = attr
        self._warmup_start = warmup_start
        self._warmup_iter = warmup_iter
        self._init = init
        self._optimizer = optimizer
        self._t = 0

    def initialize(self, trainer):
        optimizer = self._optimizer or trainer.updater.get_optimizer('main')
        if self._warmup_start is None:
            self._warmup_start = getattr(optimizer, self._attr)
        else:
            setattr(optimizer, self._attr, self._warmup_start)

    def __call__(self, trainer):
        self._t += 1
        if self._t <= self._warmup_iter:
            optimizer = self._optimizer or \
                trainer.updater.get_optimizer('main')
            value = (self._t * self._init + (self._warmup_iter - self._t)
                     * self._warmup_start) / self._warmup_iter
            setattr(optimizer, self._attr, value)

    def serialize(self, serializer):
        self._t = serializer('_t', self._t)
