from torch.optim.lr_scheduler import CosineAnnealingLR, _LRScheduler


class DelayerScheduler(_LRScheduler):
    """Starts with a flat lr schedule until it reaches N epochs the applies a scheduler
    Args:
            optimizer (Optimizer): Wrapped optimizer.
            delay_epochs: number of epochs to keep the initial lr until starting aplying the scheduler
            after_scheduler: after target_epoch, use this scheduler(eg. ReduceLROnPlateau)
    """

    def __init__(self, optimizer, delay_epochs, after_scheduler):
        self.delay_epochs = delay_epochs
        self.after_scheduler = after_scheduler
        self.finished = False
        super().__init__(optimizer)

    def get_lr(self):
        if self.last_epoch >= self.delay_epochs:
            if not self.finished:
                self.after_scheduler.base_lrs = self.base_lrs
                self.finished = True
            return self.after_scheduler.get_lr()

        return self.base_lrs

    def step(self, epoch=None):
        if self.finished:
            if epoch is None:
                self.after_scheduler.step(None)
            else:
                self.after_scheduler.step(epoch - self.delay_epochs)
        else:
            return super(DelayerScheduler, self).step(epoch)


def DelayedCosineAnnealingLR(optimizer, delay_epochs, cosine_annealing_epochs):
    base_scheduler = CosineAnnealingLR(optimizer, cosine_annealing_epochs)
    return DelayerScheduler(optimizer, delay_epochs, base_scheduler)
