import math


def step_decay(lr0, s, epochs_drop=1.0):
    """
    Create stepwise decay: Drop learning rate by half (s) every specified iteration.

    Parameters
    ----------
    lr0 : float
        initial learning rate
    s: float
        decay rate, e.g. 0.5, choose lower s than for other decays
    epochs_drop: float
        step size

    Returns
    -------
    step_decay_fn: float
        stepwise decay
    """
    # initial_lrate = 0.1
    # drop = 0.5
    # epochs_drop = 1.0
    def step_decay_fn(steps_per_epoch):
        return lr0 * math.pow(s, math.floor((1 + steps_per_epoch) / epochs_drop))
    return step_decay_fn