import math


def get_learning_rate(lr_schedule,
                      tot_epoch,
                      cur_epoch):
    """ Compuate learning rate based on the initialized learning schedule,
    e.g., cosine, 1.5e-4,90,6e-5,150,0 -> ('cosine', [1.5e-4,90,6e-5,150,0]).

    Args:
        lr_schedule: the learning schedule specified in args.py
        tot_epoch: the total epoch
        cur_epoch: the current (float) epoch
    """

    # cosine
    lr_mtd = lr_schedule[0]
    # 90, 150
    ep_list = lr_schedule[1][1::2]
    # 1.5e-4,6e-5,0
    lr_list = lr_schedule[1][::2]
    assert len(ep_list) + 1 == len(lr_list)
    for start, end, lr in zip([0] + ep_list,
                              ep_list + [tot_epoch],
                              lr_list):
        if start <= cur_epoch < end:
            if lr_mtd == 'cosine':
                return lr * (math.cos((cur_epoch - start) / (end - start) * math.pi) + 1) / 2
            elif lr_mtd == 'const':
                return lr
            else:
                raise TypeError('Invalid learning method {}'.format(lr_mtd))
    raise TypeError('Invalid learning schedule {}'.format(lr_schedule))