import math


def lr_schedule(step, total_steps, lr, lr_warmup, lr_cooldown):
    """Calculate a scheduled learning rate value."""
    if lr_warmup is None:
        lr_warmup = 0
    else:
        if 0 < lr_warmup <= 1:
            lr_warmup = math.ceil(lr_warmup * total_steps)
        if step <= lr_warmup:
            return lr * (step / lr_warmup)
    if lr_cooldown == 'linear':
        return lr * (1.0 - (step - lr_warmup) /
                           (total_steps - lr_warmup))
    else:
        return lr