import math


def binary_eps_search(eps_lower_bound, eps_upper_bound, bab_function, quantization=1e-3, mode="LB"):
    """
    Run binary search on the epsilon values in order to create a BaB dataset.
    :parameter eps_lower_bound: starting lower bound on epsilon
    :parameter eps_upper_bound: starting upper bound on epsilon
    :parameter bab_function: BaB function, takes only epsilon as input
    :parameter quantization: how spaced apart are the epsilons we are considering (default 1e-3).
        The effective quantization is the largest (eps_upper_bound - eps_lower_bound)/2^k > quantization.
        Min quantization/2, max quantization.

    ### Search criterion:
    LB mode: the property with min-eps (within quantization) that is either SAT or has timed out.
    UB mode: the property with max-eps (within quantization) that is either UNSAT or has timed out.

    Returns result, rounded upwards using the quantization.
    """
    assert mode in ["LB", "UB"]
    print(f"Starting epsilon bounds: LB: {eps_lower_bound}, UB: {eps_upper_bound}")

    while (eps_upper_bound - eps_lower_bound) > quantization:
        c_epsilon = (eps_upper_bound + eps_lower_bound) / 2

        # Run BaB with the current epsilon value.
        bab_status, bab_runtime = bab_function(c_epsilon)
        print(f"BaB status {bab_status}, BaB runtime {bab_runtime}")

        conditions = ["True"] if mode == "UB" else ["True", "timeout"]
        if bab_status in conditions:
            eps_upper_bound = c_epsilon
        else:
            eps_lower_bound = c_epsilon

        print(f"Current epsilon bounds: LB: {eps_lower_bound}, UB: {eps_upper_bound}")

    return_value = math.floor(eps_lower_bound / quantization) * quantization if mode == "UB" else \
        math.ceil(eps_upper_bound / quantization) * quantization
    return return_value