"""PyTorch edition of TensorFlow learning schedule in tensorflow object
detection API.
"""
import numpy as np
from torch.optim.optimizer import Optimizer


class _LRSchedulerStep(object):
    def __init__(self, optimizer, last_step=-1):
        if not isinstance(optimizer, Optimizer):
            raise TypeError(
                "{} is not an Optimizer".format(type(optimizer).__name__)
            )
        self.optimizer = optimizer
        if last_step == -1:
            for group in optimizer.param_groups:
                group.setdefault("initial_lr", group["lr"])
        else:
            for i, group in enumerate(optimizer.param_groups):
                if "initial_lr" not in group:
                    raise KeyError(
                        "param 'initial_lr' is not specified "
                        "in param_groups[{}] when resuming an optimizer".format(
                            i
                        )
                    )
        self.base_lrs = list(
            map(lambda group: group["initial_lr"], optimizer.param_groups)
        )
        self.step(last_step + 1)
        self.last_step = last_step

    """
    def get_lr(self):
        raise NotImplementedError
    """

    def get_lr(self):
        ret = [self._get_lr_per_group(base_lr) for base_lr in self.base_lrs]
        return ret

    def _get_lr_per_group(self, base_lr):
        raise NotImplementedError

    def step(self, step=None):
        if step is None:
            step = self.last_step + 1
        self.last_step = step
        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
            param_group["lr"] = lr


class Constant(_LRSchedulerStep):
    def __init__(self, optimizer, last_step=-1):
        super().__init__(optimizer, last_step)

    def _get_lr_per_group(self, base_lr):
        return base_lr


class ManualStepping(_LRSchedulerStep):
    """Pytorch edition of manual_stepping in tensorflow.
    DON'T SUPPORT PARAM GROUPS.
    """

    def __init__(self, optimizer, boundaries, rates, last_step=-1):
        self._boundaries = boundaries
        self._num_boundaries = len(boundaries)
        self._learning_rates = rates

        if any([b < 0 for b in boundaries]) or any(
            [not isinstance(b, int) for b in boundaries]
        ):
            raise ValueError("boundaries must be a list of positive integers")
        if any(
            [bnext <= b for bnext, b in zip(boundaries[1:], boundaries[:-1])]
        ):
            raise ValueError(
                "Entries in boundaries must be strictly increasing."
            )
        if any([not isinstance(r, float) for r in rates]):
            raise ValueError("Learning rates must be floats")
        if len(rates) != len(boundaries) + 1:
            raise ValueError(
                "Number of provided learning rates must exceed "
                "number of boundary points by exactly 1."
            )
        super().__init__(optimizer, last_step)

    def _get_lr_per_group(self, base_lr):
        step = self.last_step
        ret = None
        for i, bound in enumerate(self._boundaries):
            if step > bound:
                ret = self._learning_rates[i + 1]
        if ret is not None:
            return ret
        return self._learning_rates[0]


class ExponentialDecayWithBurnin(_LRSchedulerStep):
    """Pytorch edition of manual_stepping in tensorflow.
    """

    def __init__(
        self,
        optimizer,
        learning_rate_decay_steps,
        learning_rate_decay_factor,
        burnin_learning_rate,
        burnin_steps,
        last_step=-1,
    ):
        self._decay_steps = learning_rate_decay_steps
        self._decay_factor = learning_rate_decay_factor
        self._burnin_learning_rate = burnin_learning_rate
        self._burnin_steps = burnin_steps

        super().__init__(optimizer, last_step)

    def _get_lr_per_group(self, base_lr):
        if self._burnin_learning_rate == 0:
            burnin_learning_rate = base_lr
        step = self.last_step
        post_burnin_learning_rate = base_lr * self._decay_factor ^ (
            step // self._decay_steps
        )
        if step < self._burnin_steps:
            return burnin_learning_rate
        else:
            return post_burnin_learning_rate


class ExponentialDecay(_LRSchedulerStep):
    def __init__(
        self,
        optimizer,
        learning_rate_decay_steps,
        learning_rate_decay_factor,
        staircase=True,
        last_step=-1,
    ):
        self._decay_steps = learning_rate_decay_steps
        self._decay_factor = learning_rate_decay_factor
        self._staircase = staircase

        super().__init__(optimizer, last_step)

    def _get_lr_per_group(self, base_lr):
        step = self.last_step
        if self._staircase:
            post_burnin_learning_rate = base_lr * pow(
                self._decay_factor, (step // self._decay_steps)
            )
        else:
            post_burnin_learning_rate = base_lr * pow(
                self._decay_factor, (step / self._decay_steps)
            )

        return post_burnin_learning_rate


class CosineDecayWithWarmup(_LRSchedulerStep):
    def __init__(
        self,
        optimizer,
        total_steps,
        warmup_learning_rate,
        warmup_steps,
        last_step=-1,
    ):
        if total_steps < warmup_steps:
            raise ValueError(
                "total_steps must be larger or equal to " "warmup_steps."
            )
        self._total_steps = total_steps
        self._warmup_learning_rate = warmup_learning_rate
        self._warmup_steps = warmup_steps

        super().__init__(optimizer, last_step)

    def _get_lr_per_group(self, base_lr):
        if base_lr < self._warmup_learning_rate:
            raise ValueError(
                "learning_rate_base must be larger "
                "or equal to warmup_learning_rate."
            )

        step = self.last_step
        learning_rate = (
            0.5 *
            base_lr *
            (
                1 +
                np.cos(
                    np.pi *
                    (float(step) - self._warmup_steps) /
                    float(self._total_steps - self._warmup_steps)
                )
            )
        )
        if self._warmup_steps > 0:
            slope = (base_lr - self._warmup_learning_rate) / self._warmup_steps
            pre_cosine_learning_rate = (
                slope * float(step) + self._warmup_learning_rate
            )
            if step < self._warmup_steps:
                return pre_cosine_learning_rate
            else:
                return learning_rate
