import math


def get_learning_rate(iter, num_jobs, num_iters, num_archives_processed,
                      num_archives_to_process,
                      initial_effective_lrate, final_effective_lrate,
                      schedule_type='linear'):
    """Get learning rate for current iteration based on the learning rate schedule.

    This function implements 3 learning rate schedules (see description of 'schedule_type'
    argument for more information). It returns the learning based on the schedule selected.
    This function is stateless, so we don't assume that the entire learning process
    uses the same schedule.

    Args:
        iter (int): current iteration
        num_jobs (int): number of jobs to be run for this iteration
        num_iters (int): total number of iters
        num_archives_processed (int): number of archives processed till now
        num_archives_to_process (int): number of archives to be processed
        initial_effective_lrate (float): initial learning rate
        final_effective_lrate (float): final learning rate
        schedule_type (str): Default value is 'linear'. The function supports 3 modes
            1. 'none': no change in learning rate. the initial_effective_lrate will be used
            2. 'linear' (default): reduce the learning rate linearly
            3. 'exponential': this is what Kaldi implements. the learning rate reduces exponentially.

    Returns:
        Effective learning rate: learning rate based on the schedule multiplied by the number
            of jobs based on the mode (for linear and exponential only)
    """
    if schedule_type == 'none':
        return initial_effective_lrate
    elif schedule_type == 'linear':
        epoch_no = (num_archives_processed // num_archives_to_process) + 1
        return (initial_effective_lrate/epoch_no)*num_jobs
    elif schedule_type == 'exponential':
        if iter + 1 >= num_iters:
            effective_learning_rate = final_effective_lrate
        else:
            effective_learning_rate = (
                    initial_effective_lrate
                    * math.exp(num_archives_processed
                               * math.log(float(final_effective_lrate) / initial_effective_lrate)
                               / num_archives_to_process))

        return num_jobs * effective_learning_rate