import math


def cyclical_annealing(train_step: int, cycle_size: int, cycle_size_const_fraction: float = 0.5,
                       beta_final: float = 1.0, beta_start: float = 0.0) -> float:
    """Calculate beta (KL term weight) as a function of the train step with cyclic annealing.

    How to calculate the cycle size:

    The cyclic scheduling scheme looks like this. The "---" part is cycle_size_const_fraction from the
    whole cycle_size where beta = beta_final. Before a monotonic beta increase happens from beta_start to beta_final.
    beta_final           ---  ---  ---
                        /  | /  | /  |
    beta_start         /   |/   |/   | ...
                      |----|
                    cycle_size
    """
    first_const_step = math.ceil(cycle_size - cycle_size * cycle_size_const_fraction)  # within a cycle
    step_in_cycle = train_step % cycle_size
    if step_in_cycle < first_const_step:
        slope = (beta_final - beta_start) / first_const_step
        return slope * step_in_cycle + beta_start
    else:
        return beta_final