:- module bug155b.

:- interface.

:- import_module float.
:- import_module list.
:- import_module set.

:- type best_solutions(T)
    --->    no_best_solutions
    ;       best_solutions(
                bs_solutions            :: list(T),
                bs_objective_value      :: float
                    % Note that the solver tries to minimise this value.  That
                    % is smaller numbers are more optimal.
            ).

:- type bnb_state(T).

:- pred branch_and_bound(semipure pred(bnb_state(T), T),
    func(T) = float, set(T)).
:- mode branch_and_bound(pred(in, out) is nondet, 
    func(in) = out is det, out) is det.

:- implementation.

:- import_module io.
:- import_module mutvar.
:- import_module string.
:- import_module require.

:- type bnb_state(T)
    --->    bnb_state(
                best_solutions_mutable      :: mutvar(best_solutions(T)),
                objective_function          :: func(T) = float
            ).

:- inst bnb_state
    --->    bnb_state(ground, func(in) = out is det).

branch_and_bound(Generator, ObjectiveFn, FinalBestSolutions) :-
    % Use a failure driven loop to implement a branch and bound solver.
    promise_pure (
        impure new_mutvar(no_best_solutions, BestSolutionsMutvar),
        State = bnb_state(BestSolutionsMutvar, ObjectiveFn),
        (
            semipure Generator(State, CurrentSolution),

            CurrentObjective = ObjectiveFn(CurrentSolution),
           
            impure get_mutvar(BestSolutionsMutvar, BestSolutions0),
            (
                BestSolutions0 = no_best_solutions,
                BestSolutions = best_solutions([CurrentSolution],
                    CurrentObjective)
            ;
                BestSolutions0 = best_solutions(Solutions, BestObjective),
                ( CurrentObjective < BestObjective ->
                    BestSolutions = best_solutions([CurrentSolution],
                        CurrentObjective)
                ; CurrentObjective = BestObjective ->
                    BestSolutions = best_solutions(
                        [CurrentSolution | Solutions], BestObjective)
                ;
                    fail
                )
            ),
            impure set_mutvar(BestSolutionsMutvar, BestSolutions),
            semidet_fail
        ->
            error("Failure driven loop must fail")
        ;
            impure get_mutvar(BestSolutionsMutvar, FinalBestSolutions0)
        ),
        (  
            FinalBestSolutions0 = no_best_solutions,
            FinalBestSolutions = set.init
        ;  
            FinalBestSolutions0 = best_solutions(Solutions, _),
            FinalBestSolutions = set.from_list(Solutions)
        )
    ).