from torch.optim.lr_scheduler import _LRScheduler


# This is rather suboptimal, because we need to import a protected class. Unfortunately, I don't see another way.


class ToucanWarmupScheduler(_LRScheduler):
    """
    A warmup scheduler that should be called after every batch.
    """

    def __init__(self, optimizer, peak_lr=0.0002, warmup_steps=20000, max_steps=200000, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.peak_lr = peak_lr
        self.max_steps = max_steps
        self.plateau = self.warmup_steps * 4
        self.last_lr = 0.0
        # __init__() must be invoked before setting field
        # because step() is also invoked in __init__()
        super().__init__(optimizer, last_epoch)

    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps})"

    def get_lr(self):
        step_num = self.last_epoch + 1
        if step_num <= self.warmup_steps:
            lr = self.peak_lr * min(step_num / self.warmup_steps, 1.0)
            self.last_lr = lr
            return [lr for _ in self.base_lrs]
        elif step_num < self.warmup_steps + self.plateau:
            self.last_lr = self.peak_lr
            return [self.peak_lr for _ in self.base_lrs]
        else:
            scale = 1 - (((step_num - (self.warmup_steps + self.plateau)) / self.max_steps) / (self.max_steps / 10))
            self.last_lr = max(self.last_lr * scale, 1e-7)
            return [self.last_lr for _ in self.base_lrs]


class WarmupScheduler(_LRScheduler):
    """
    The WarmupLR scheduler
    This scheduler is almost same as NoamLR Scheduler except for following difference:
    NoamLR:
        lr = optimizer.lr * model_size ** -0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)
    WarmupLR:
        lr = optimizer.lr * warmup_step ** 0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)
    Note that the maximum lr equals to optimizer.lr in this scheduler.

    Taken from ESPnet
    """

    def __init__(self, optimizer, warmup_steps=25000, last_epoch=-1):
        self.warmup_steps = warmup_steps
        # __init__() must be invoked before setting field
        # because step() is also invoked in __init__()
        super().__init__(optimizer, last_epoch)

    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps})"

    def get_lr(self):
        step_num = self.last_epoch + 1
        return [lr * self.warmup_steps ** 0.5 * min(step_num ** -0.5, step_num * self.warmup_steps ** -1.5) for lr in
                self.base_lrs]


if __name__ == '__main__':
    lrs = list()
    warmup_steps = 30000
    peak_lr = 0.0005
    max_steps = 800000
    plateau_size = warmup_steps * 5
    for step_num in range(max_steps):
        if step_num <= warmup_steps:
            lr = peak_lr * min(step_num / warmup_steps, 1.0)
            lrs.append(lr)
        elif step_num < warmup_steps + plateau_size:
            lrs.append(peak_lr)
        else:
            scale = 1 - (((step_num - (warmup_steps + plateau_size)) / max_steps) / (max_steps / 10))
            lrs.append(max(lrs[-1] * scale, 1e-7))
    import matplotlib.pyplot as plt

    plt.plot(lrs)
    plt.show()