import enum
import numpy as np


def inv_lr_scheduler(param_lr, optimizer, iter_num, gamma=0.0001,
                     power=0.75, init_lr=0.001):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
    lr = init_lr * (1 + gamma * iter_num) ** (- power)
    i = 0
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr * param_lr[i]
        i += 1
    return optimizer


def calc_coeff(iter_num, high=1.0, low=0.0, alpha=10.0, max_iter=10000.0):
    return np.float(2.0 * (high - low) /
                    (1.0 + np.exp(- alpha * iter_num / max_iter)) -
                    (high - low) + low)


class InvLr:
    def __init__(self, optimizer, gamma=0.0001, power=0.75):
        self.optimizer = optimizer
        self.gamma = gamma
        self.power = power
        self.init_lr = []
        for param_group in optimizer.param_groups:
            self.init_lr.append(param_group['lr'])

    def step(self, step):
        factor = (1 + self.gamma * step) ** (- self.power)
        for index, param_group in enumerate(self.optimizer.param_groups):
            param_group['lr'] = self.init_lr[index] * factor
