from typing import Optional
from torch.optim.optimizer import Optimizer

class StepwiseLR:
    """
    A lr_scheduler that update learning rate using the following schedule:

    .. math::
        \text{lr} = \text{init_lr} \times \text{lr_mult} \times (1+\gamma i)^{-p},

    where `i` is the iteration steps.

    Parameters:
        - **optimizer**: Optimizer
        - **init_lr** (float, optional): initial learning rate. Default: 0.01
        - **gamma** (float, optional): :math:`\gamma`. Default: 0.001
        - **decay_rate** (float, optional): :math:`p` . Default: 0.75
    """
    def __init__(self, optimizer: Optimizer, init_lr: Optional[float] = 0.01,
                 gamma: Optional[float] = 0.001, decay_rate: Optional[float] = 0.75):
        self.init_lr = init_lr
        self.gamma = gamma
        self.decay_rate = decay_rate
        self.optimizer = optimizer
        self.iter_num = 0

    def get_lr(self) -> float:
        lr = self.init_lr * (1 + self.gamma * self.iter_num) ** (-self.decay_rate)
        return lr

    def step(self):
        # """Increase iteration number `i` by 1 and update learning rate in `optimizer`"""
        # lr = self.get_lr()
        # for param_group in self.optimizer.param_groups:
        #     if 'lr_mult' not in param_group:
        #         param_group['lr_mult'] = 1.
        #     param_group['lr'] = lr * param_group['lr_mult']
        #
        # self.iter_num += 1

        lr = self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr * param_group['lr_mult']

        self.iter_num+=1



