def __update_agent_state(current_agent_state, transition_probability, rng):
    """
    Get agent state for next time step

    Parameters
    ----------
    current_agent_state : int
        Current agent state.
    transition_probability : ndarray
        Transition probability vector corresponding to current state.
    rng : numpy random number generator

    Returns
    -------
    agent_state : int
        Agent state at next time step.

    """
    choice = rng.uniform()

    agent_state = current_agent_state

    for i in range(len(transition_probability)):
        if choice < sum(transition_probability[0: i + 1]):
            agent_state = i
            break

    return agent_state