hercules 0.5.0

use crate::qubo::Qubo;
use ndarray::Array1;
use rayon::prelude::*;

use crate::branch_node::QuboBBNode;
use crate::branch_stratagy::{BranchResult, BranchStrategy};
use crate::branch_subproblem::{get_sub_problem_solver, SubProblemSolver};
use crate::branchbound_utils::{check_integer_feasibility, get_current_time};
use crate::branchboundlogger::SolverOutputLogger;
use crate::lower_bound::li_lower_bound;
use crate::preprocess;
use crate::preprocess::preprocess_qubo;
use crate::solver_options::SolverOptions;
use std::collections::BinaryHeap;

/// Struct for the B&B Solver
pub struct BBSolver {
    pub qubo: Qubo,
    pub qubo_pp_form: Qubo,
    pub best_solution: Array1<usize>,
    pub best_solution_value: f64,
    pub nodes: BinaryHeap<QuboBBNode>,
    pub nodes_processed: usize,
    pub nodes_solved: usize,
    pub nodes_visited: usize,
    pub time_start: f64,
    pub branch_strategy: BranchStrategy,
    pub subproblem_solver: Box<dyn SubProblemSolver + Sync>,
    pub options: SolverOptions,
    pub early_stop: bool,
    pub solver_logger: SolverOutputLogger,
}

pub enum Event {
    UpdateBestSolution(Array1<usize>, f64),
    AddBranches(QuboBBNode, QuboBBNode),
    Nill,
}

pub enum NodeLoggingAction {
    Visited,
    Processed,
    Solved,
}

pub enum PruneAction {
    Prune,
    Dont,
}

pub struct ProcessNodeState {
    pub prune_action: PruneAction,
    pub events: Vec<Event>,
    pub logging: NodeLoggingAction,
}

pub enum SolverResult {
    OptimalSolution(Array1<f64>, f64),
    SubOptimalSolution(Array1<f64>, f64),
}

impl BBSolver {
    /// Creates a new B&B solver
    pub fn new(qubo: Qubo, options: SolverOptions) -> Self {

        // make sure the QUBO is in symmetric form
        let qubo = qubo.convex_symmetric_form();

        // create auxiliary variables
        let num_x = qubo.num_x();

        let subproblem_solver = get_sub_problem_solver(&qubo, &options.sub_problem_solver);
        let branch_strategy = options.branch_strategy;
        let start_time = get_current_time();
        let output_level = options.verbose;
        let pp_form = preprocess::shift_qubo(&qubo);

        Self {
            qubo,
            qubo_pp_form: pp_form,
            best_solution: Array1::zeros(num_x),
            best_solution_value: 0.0,
            nodes: BinaryHeap::new(),
            nodes_processed: 0,
            nodes_visited: 0,
            nodes_solved: 0,
            time_start: start_time,
            branch_strategy,
            subproblem_solver,
            options,
            early_stop: false,
            solver_logger: SolverOutputLogger::new(output_level),
        }
    }

    /// This function is used to warm start the solver with an initial solution if one is not provided
    pub fn warm_start(&mut self, initial_solution: Array1<usize>) {
        let warm_start_value = self.qubo.eval_usize(&initial_solution);
        self.update_solution_if_better(&initial_solution, warm_start_value);
    }

    /// The main solve function of the B&B algorithm
    pub fn solve(&mut self) -> (Array1<usize>, f64) {
        // preprocess the problem
        let fixed_variables =
            preprocess_qubo(&self.qubo_pp_form, &self.options.fixed_variables, true);
        self.options.fixed_variables.clone_from(&fixed_variables);

        // create the root node
        let mut root_node = QuboBBNode {
            lower_bound: f64::NEG_INFINITY,
            solution: 0.5 * Array1::ones(self.qubo.num_x()), // initial guess is 0.5 for all variables
            fixed_variables,
        };

        // solve the root node subproblem
        let (root_lower_bound, root_solution) = self.solve_node(&root_node);
        root_node.lower_bound = root_lower_bound;
        root_node.solution = root_solution;

        // add the root node to the list of nodes
        self.nodes.push(root_node);

        // Reset start time as it can be different from the time we created the solver instance
        self.time_start = get_current_time();

        // set up the output of the solver
        // display the header
        self.solver_logger.output_header(self);

        // if the best solution is negative, then we output the warm start information
        if self.best_solution_value < 0.0 {
            self.solver_logger.output_warm_start_info(self);
        }

        // until we have hit a termination condition, we will keep iterating
        while !(*self).termination_condition() {
            // get the most recent K nodes to process
            let nodes = self.get_next_nodes(self.options.threads);

            let process_results = nodes
                .par_iter()
                .map(|node| self.process_node(node))
                .collect::<Vec<_>>();

            // apply all the events from the parallel loop back to the solver
            for state in process_results {
                self.apply_events(state.events);
                self.apply_logging_action(state.logging);
            }

            // display the line, if verbose
            self.solver_logger.generate_output_line(self);
        }

        // display the exit line
        self.solver_logger.generate_exit_line(self);

        (self.best_solution.clone(), self.best_solution_value)
    }

    /// Checks if we can prune the node, based on the lower bound and best solution, returns an action
    pub fn can_prune_action(&self, node: &QuboBBNode) -> (PruneAction, Event) {
        // if our parent solution is above our current feasible soltion then prune
        if node.lower_bound > self.best_solution_value {
            return (PruneAction::Prune, Event::Nill);
        }

        // if the solution is complete, then we can update the best solution if better
        // we can also prune the node, as there are no more variables to fix
        if node.fixed_variables.len() == self.qubo.num_x() {
            // generate the solution vector
            let mut solution = Array1::zeros(self.qubo.num_x());
            for (&index, &value) in &node.fixed_variables {
                solution[index] = value;
            }

            let value = self.qubo.eval_usize(&solution);
            // evaluate the solution against the best solution we have so far
            // if we have a better solution update it
            return (
                PruneAction::Prune,
                Event::UpdateBestSolution(solution, value),
            );
        }

        (PruneAction::Dont, Event::Nill)
    }

    // apply the logging action to the solver
    pub fn apply_logging_action(&mut self, action: NodeLoggingAction) {
        match action {
            NodeLoggingAction::Visited => {
                // increment the number of nodes visited
                self.nodes_visited += 1;
            }
            NodeLoggingAction::Processed => {
                // increment the number of nodes processed
                self.nodes_processed += 1;
            }
            NodeLoggingAction::Solved => {
                // increment the number of nodes solved and processed
                self.nodes_processed += 1;
                self.nodes_solved += 1;
            }
        }
    }

    /// main loop of the branch and bound algorithm
    pub fn process_node(&self, node: &QuboBBNode) -> ProcessNodeState {
        // create a mutable copy of the node
        let mut node = node.clone();

        // pass to the presolver to see if there are any variables we can fix
        node.fixed_variables = preprocess_qubo(&self.qubo_pp_form, &node.fixed_variables, true);

        // calculate the lower bound via the li lower bound formula
        let li_bound = li_lower_bound(&self.qubo, &node.fixed_variables);
        node.lower_bound = node.lower_bound.max(li_bound);

        // with this expanded set, can we prune the node?
        let (prune_action, event) = self.can_prune_action(&node);

        // if we are pruning at this stage, then we can early return
        if matches!(prune_action, PruneAction::Prune) {
            return ProcessNodeState {
                prune_action,
                events: vec![event],
                logging: NodeLoggingAction::Processed,
            };
        }

        // check if integer-feasible solution
        // if not all variables are fixed, we can still check if we are 'near' integer-feasible (within 1E-10) of 0 or 1
        let (is_int_feasible, rounded_sol) = check_integer_feasibility(&node);

        // if we are integer-feasible, then we can prune this branch and return the solution
        if is_int_feasible {
            // compute the objective
            let value = self.qubo.eval_usize(&rounded_sol);

            // we will attempt to update the solution otherwise prune
            return ProcessNodeState {
                prune_action,
                events: vec![Event::UpdateBestSolution(rounded_sol, value)],
                logging: NodeLoggingAction::Solved,
            };
        }

        // We now need to solve the node to generate the lower bound and solution
        let (lower_bound, solution) = self.solve_node(&node);

        // inject the solution back into the node
        node.solution.clone_from(&solution);

        // determine what variable we are branching on
        let branch_result = self.make_branch(&node);

        // we now apply the new fixed variables to the base node before we branch
        for (&index, &value) in &branch_result.found_fixed_vars {
            node.fixed_variables.insert(index, value);
        }

        // if we have now fixed all variables, we can check if we have a solution
        if node.fixed_variables.len() == self.qubo.num_x() {
            // generate the solution vector
            let mut solution = Array1::zeros(self.qubo.num_x());
            for (&index, &value) in &node.fixed_variables {
                solution[index] = value;
            }

            let value = self.qubo.eval_usize(&solution);
            // evaluate the solution against the best solution we have so far
            // if we have a better solution update it
            return ProcessNodeState {
                prune_action,
                events: vec![Event::UpdateBestSolution(solution, value)],
                logging: NodeLoggingAction::Solved,
            };
        }

        // if we are going to branch, then we can generate a heuristic solution
        let (heur_sol, heur_obj) = self.options.heuristic.make_heuristic(self, &node);

        // generate the branches
        let (zero_branch, one_branch) =
            Self::branch(node, branch_result.branch_variable, lower_bound, solution);

        ProcessNodeState {
            prune_action,
            events: vec![
                Event::AddBranches(zero_branch, one_branch),
                Event::UpdateBestSolution(heur_sol, heur_obj),
            ],
            logging: NodeLoggingAction::Solved,
        }
    }

    pub fn apply_events(&mut self, events: Vec<Event>) {
        for action in events {
            match action {
                Event::UpdateBestSolution(solution, value) => {
                    self.update_solution_if_better(&solution, value);
                }
                Event::AddBranches(zero_branch, one_branch) => {
                    // only add the branches if their lower bound is better than the current best solution
                    if zero_branch.lower_bound <= self.best_solution_value {
                        self.nodes.push(zero_branch);
                    }
                    if one_branch.lower_bound <= self.best_solution_value {
                        self.nodes.push(one_branch);
                    }
                }
                Event::Nill => {}
            }
        }
    }

    /// update the best solution if better than the current best solution
    pub fn update_solution_if_better(&mut self, solution: &Array1<usize>, solution_value: f64) {
        if solution_value < self.best_solution_value {
            self.best_solution.clone_from(solution);
            self.best_solution_value = solution_value;

            // if we have an early stopping condition, then we can check if we have a solution
            // let beck_proof = beck_proof(&self.qubo, &self.best_solution);
            //
            // // if we have a beck proof, then we can stop early
            // if beck_proof {
            //     self.early_stop = true;
            //     self.solver_logger.early_termination();
            // }

            // We can remove all nodes that are worse than the current best solution
            self.nodes.retain(|node| {
                // if the node's lower bound is worse than the best solution, we can prune it
                node.lower_bound <= self.best_solution_value
            });
        }
    }

    /// This function is used to get the next node to process, popping it from the list of nodes
    pub fn get_next_node(&mut self) -> Option<QuboBBNode> {
        while !self.nodes.is_empty() {
            // we pull a node from our node list
            let optional_node = self.nodes.pop();

            // check and unwrap the node if it is safe
            let node = optional_node?;

            // we increment the number of nodes we have visited
            self.apply_logging_action(NodeLoggingAction::Visited);

            // if we can't prune it, then we return it
            let (prune, event) = self.can_prune_action(&node);

            // if we have stumbled into a better solution, then we can take it
            if let Event::UpdateBestSolution(solution, value) = event {
                self.update_solution_if_better(&solution, value);
            }

            // if we don't prune the node then we can return it
            if matches!(prune, PruneAction::Dont) {
                return Some(node);
            }
        }

        None
    }

    pub fn get_next_nodes(&mut self, n: usize) -> Vec<QuboBBNode> {
        let mut nodes = Vec::new();

        // loop while we haven't filled our vector OR the node list is not empty
        while nodes.len() <= n {
            let next_node = self.get_next_node();

            // if there is a node to add, do so, else break out as there aren't any nodes left
            if let Some(node) = next_node {
                nodes.push(node);
            } else {
                break;
            }
        }

        nodes
    }

    /// Checks for termination conditions of the B&B algorithm, such as time limit or no more nodes
    pub fn termination_condition(&self) -> bool {
        // get current time to check if we have exceeded the maximum time
        let current_time = get_current_time();

        // check if we violated the time limit
        if current_time - self.time_start > self.options.max_time {
            return true;
        }

        // check if we have no more nodes to process
        if self.nodes.is_empty() {
            return true;
        }

        // if we have an early stopping condition, then we can check if we have a solution
        if self.early_stop {
            return true;
        }

        false
    }

    /// Branch Selection Strategy - Currently selects the first variable that is not fixed
    pub fn make_branch(&self, node: &QuboBBNode) -> BranchResult {
        self.branch_strategy.make_branch(self, node)
    }

    /// Actually branches the node into two new nodes
    pub fn branch(
        node: QuboBBNode,
        branch_id: usize,
        lower_bound: f64,
        solution: Array1<f64>,
    ) -> (QuboBBNode, QuboBBNode) {
        // make two new nodes that are clones of the parent, one with the variable set to 0 and
        // the other set to 1
        let mut zero_branch = node.clone();
        let mut one_branch = node;

        // add fixed variables
        zero_branch.fixed_variables.insert(branch_id, 0);
        one_branch.fixed_variables.insert(branch_id, 1);

        // update the solution and lower bound for the new nodes
        zero_branch.solution.clone_from(&solution);
        one_branch.solution = solution;

        // set the lower bound for the new nodes
        zero_branch.lower_bound = lower_bound;
        one_branch.lower_bound = lower_bound;

        (zero_branch, one_branch)
    }

    pub fn solve_node(&self, node: &QuboBBNode) -> (f64, Array1<f64>) {
        self.subproblem_solver.solve_lower_bound(self, node, None)
    }
}

#[cfg(test)]
mod tests {
    use crate::branch_stratagy::BranchStrategy;
    use crate::branch_subproblem::SubProblemSelection;
    use crate::preprocess::preprocess_qubo;
    use crate::qubo::Qubo;
    use crate::solver_options::SolverOptions;
    use crate::tests::make_test_prng;
    use crate::{branchbound, local_search};
    use ndarray::Array1;
    use sprs::CsMat;
    use std::collections::HashMap;

    pub fn get_default_solver_options() -> SolverOptions {
        let mut options = SolverOptions::new();
        options.verbose = 1;
        options.max_time = 1000.0;
        options.threads = 20;
        options
    }

    #[test]
    pub fn branch_bound_test() {
        let mut prng = make_test_prng();
        let eye = CsMat::eye(3);
        let c = Array1::from_vec(vec![-1.1, -2.0, -3.0]);
        let p = Qubo::new_with_c(eye, c);

        let guess = local_search::particle_swarm_search(&p, 100, 1000, &mut prng);
        let mut solver = branchbound::BBSolver::new(p, SolverOptions::new());
        solver.warm_start(guess);
        solver.solve();

        assert_eq!(solver.best_solution, Array1::from_vec(vec![1, 1, 1]));
    }

    #[test]
    pub fn test_gka2b_solve() {
        let file_path = "test_data/gka2b.qubo";
        let p = Qubo::read_qubo(file_path);

        let sol_val = Array1::from_vec(vec![
            0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0,
            0, 0,
        ]);

        solve_qubo_with_all_permutations(&p, &sol_val);
    }

    #[test]
    pub fn test_gka1b_solve() {
        let file_path = "test_data/gka1b.qubo";
        let p = Qubo::read_qubo(file_path);

        let sol_val = Array1::from_vec(vec![
            0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0,
        ]);
        solve_qubo_with_all_permutations(&p, &sol_val);
    }

    #[test]
    pub fn test_gka6a_solve() {
        let file_path = "test_data/gka6a.qubo";
        let p = Qubo::read_qubo(file_path);

        let sol_val = Array1::from_vec(vec![
            0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1,
            0, 1,
        ]);
        solve_qubo_with_all_permutations(&p, &sol_val);
    }

    #[test]
    pub fn test_gka7a_solve() {
        let file_path = "test_data/gka7a.qubo";
        let p = Qubo::read_qubo(file_path);

        let sol_val = Array1::from_vec(vec![
            0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0,
            0, 1,
        ]);
        solve_qubo_with_all_permutations(&p, &sol_val);
    }

    pub fn solve_qubo_with_all_permutations(qubo: &Qubo, sol_val: &Array1<usize>) {
        let branch_options = vec![
            BranchStrategy::FirstNotFixed,
            BranchStrategy::MostViolated,
            BranchStrategy::Random,
            BranchStrategy::WorstApproximation,
            BranchStrategy::WorstApproximation2,
            BranchStrategy::MostEdges,
            BranchStrategy::LargestEdges,
            BranchStrategy::MostFixed,
            BranchStrategy::FullStrongBranching,
            BranchStrategy::PartialStrongBranching,
            BranchStrategy::LargestDiag,
            BranchStrategy::MoveingEdges,
            BranchStrategy::RoundRobin,
        ];

        let sub_problem_solvers = vec![
            SubProblemSelection::ClarabelQP,
            SubProblemSelection::ClarabelLP,
            SubProblemSelection::HerculesCDQP,
        ];

        for branch in &branch_options {
            for sup_problem_solver in &sub_problem_solvers {
                setup_and_solve_problem(branch, &sup_problem_solver, qubo, sol_val);
            }
        }
    }

    pub fn setup_and_solve_problem(
        branch: &BranchStrategy,
        sup_problem_solver: &SubProblemSelection,
        qubo: &Qubo,
        true_sol: &Array1<usize>,
    ) {
        let mut prng = make_test_prng();

        let fixed_variables = preprocess_qubo(&qubo, &HashMap::new(), false);

        let guess = local_search::particle_swarm_search(&qubo, 10, 100, &mut prng);

        let mut options = get_default_solver_options();

        options.branch_strategy = *branch;
        options.fixed_variables = fixed_variables.clone();
        options.sub_problem_solver = *sup_problem_solver;
        options.verbose = 0;

        let mut solver = branchbound::BBSolver::new(qubo.clone(), options);

        solver.warm_start(guess);

        let (_, sol_value) = solver.solve();

        let actual_obj = solver.qubo.eval_usize(&true_sol);

        // the solution should be within 1E-5 of the actual solution
        // we don't check against the solution as there can be multiple optimal solutions
        assert!((sol_value - actual_obj).abs() <= 1E-5);
    }
}