import math

from torch.optim.lr_scheduler import _LRScheduler


class SWALR(_LRScheduler):
    def __init__(self, optimizer, swa_lr, anneal_steps=10, anneal_strategy='cos', last_epoch=-1, offset=0):
        swa_lrs = self._format_param(optimizer, swa_lr)
        for swa_lr, group in zip(swa_lrs, optimizer.param_groups):
            group['swa_lr'] = swa_lr
        if anneal_strategy not in ['cos', 'linear']:
            raise ValueError("anneal_strategy must by one of 'cos' or 'linear', "
                             "instead got {}".format(anneal_strategy))
        elif anneal_strategy == 'cos':
            self.anneal_func = self._cosine_anneal
        elif anneal_strategy == 'linear':
            self.anneal_func = self._linear_anneal
        if not isinstance(anneal_steps, int) or anneal_steps < 0:
            raise ValueError("anneal_steps must be equal or greater than 0, got {}".format(
                             anneal_steps))
        self.anneal_steps = anneal_steps
        self.offset = offset

        super(SWALR, self).__init__(optimizer, last_epoch)

    @staticmethod
    def _format_param(optimizer, swa_lrs):
        if isinstance(swa_lrs, (list, tuple)):
            if len(swa_lrs) != len(optimizer.param_groups):
                raise ValueError("swa_lr must have the same length as "
                                 "optimizer.param_groups: swa_lr has {}, "
                                 "optimizer.param_groups has {}".format(
                                     len(swa_lrs), len(optimizer.param_groups)))
            return swa_lrs
        else:
            return [swa_lrs] * len(optimizer.param_groups)

    @staticmethod
    def _linear_anneal(t):
        return t

    @staticmethod
    def _cosine_anneal(t):
        return (1 - math.cos(math.pi * t)) / 2

    def _get_initial_lr(self, lr, swa_lr, alpha):
        if self.offset > 0 and self._step_count == 1:
            return lr
        if alpha == 1:
            return swa_lr
        return (lr - alpha * swa_lr) / (1 - alpha)

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn("To get the last learning rate computed by the scheduler, "
                          "please use `get_last_lr()`.", UserWarning)
        step = self._step_count - 1
        total = max(1, self.anneal_steps + self.offset)
        if self.anneal_steps == 0:
            step = max(1, step)

        prev_t = (step - 1 + self.offset) / total
        prev_t = max(0, min(1, prev_t))
        prev_alpha = self.anneal_func(prev_t)
        prev_lrs = [self._get_initial_lr(group['lr'], group['swa_lr'], prev_alpha)
                    for group in self.optimizer.param_groups]

        t = (step + self.offset) / total
        t = max(0, min(1, t))
        alpha = self.anneal_func(t)

        lrs = [group['swa_lr'] * alpha + lr * (1 - alpha)
               for group, lr in zip(self.optimizer.param_groups, prev_lrs)]
        return lrs
