#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________

from collections import namedtuple
from heapq import heappush, heappop
import traceback

from pyomo.common.collections import ComponentMap
from pyomo.common.config import add_docstring_list
from pyomo.common.errors import InfeasibleConstraintException
from pyomo.contrib.fbbt.fbbt import fbbt
from pyomo.contrib.gdpopt.algorithm_base_class import _GDPoptAlgorithm
from pyomo.contrib.gdpopt.create_oa_subproblems import (
    add_util_block, add_disjunction_list, add_disjunct_list,
    add_algebraic_variable_list, add_boolean_variable_lists,
    add_transformed_boolean_variable_list)
from pyomo.contrib.gdpopt.config_options import (
    _add_nlp_solver_configs, _add_BB_configs, _add_mip_solver_configs,
    _add_tolerance_configs, _add_nlp_solve_configs)
from pyomo.contrib.gdpopt.nlp_initialization import (
    restore_vars_to_original_values)
from pyomo.contrib.gdpopt.util import (
    copy_var_list_values, SuppressInfeasibleWarning, get_main_elapsed_time)
from pyomo.contrib.satsolver.satsolver import satisfiable
from pyomo.core import minimize, Suffix, Constraint, TransformationFactory
from pyomo.opt import SolverFactory, SolverStatus
from pyomo.opt import TerminationCondition as tc

_linear_degrees = {1, 0}

# Data tuple for each node that also functions as the sort key.
# Therefore, ordering of the arguments below matters.
BBNodeData = namedtuple('BBNodeData', [
    'obj_lb',  # lower bound on objective value, sign corrected to minimize
    'obj_ub',  # upper bound on objective value, sign corrected to minimize
    'is_screened',  # True if the node has been screened; False if not.
    'is_evaluated',  # True if node has been evaluated; False if not.
    'num_unbranched_disjunctions',  # number of unbranched disjunctions
    'node_count',  # cumulative node counter
    'unbranched_disjunction_indices',  # list of unbranched disjunction indices
])

@SolverFactory.register(
    'gdpopt.lbb',
    doc="The LBB (logic-based branch and bound) Generalized Disjunctive "
    "Programming (GDP) solver")
class GDP_LBB_Solver(_GDPoptAlgorithm):
    """The GDPopt (Generalized Disjunctive Programming optimizer) logic-based
    branch and bound (LBB) solver.

    Accepts models that can include nonlinear, continuous variables and
    constraints, as well as logical conditions.
    """
    CONFIG = _GDPoptAlgorithm.CONFIG()
    _add_mip_solver_configs(CONFIG)
    _add_nlp_solver_configs(CONFIG, default_solver='ipopt')
    _add_nlp_solve_configs(
        CONFIG, default_nlp_init_method=restore_vars_to_original_values)
    _add_tolerance_configs(CONFIG)
    _add_BB_configs(CONFIG)

    algorithm = 'LBB'

    def _log_citation(self, config):
        config.logger.info("\n" + """- LBB algorithm:
        Lee, S; Grossmann, IE.
        New algorithms for nonlinear generalized disjunctive programming.
        Comp. and Chem. Eng. 2000, 24, 2125-2141.
        DOI: 10.1016/S0098-1354(00)00581-0.
        """.strip())

    def _solve_gdp(self, model, config):
        self.explored_nodes = 0

        # Create utility block on the original model so that we will be able to
        # copy solutions between
        util_block = self.original_util_block = add_util_block(model)
        add_disjunct_list(util_block)
        add_algebraic_variable_list(util_block)
        add_boolean_variable_lists(util_block)

        root_node = TransformationFactory(
            'contrib.logical_to_disjunctive').create_using(model)
        root_util_blk = root_node.component(
            self.original_util_block.name)
        # Add to root utility block what we will need during the algorithm
        add_disjunction_list(root_util_blk)
        # Now that logical_to_disjunctive has been called.
        add_transformed_boolean_variable_list(root_util_blk)

        # Map unfixed disjunct -> list of deactivated constraints
        root_util_blk.disjunct_to_nonlinear_constraints = ComponentMap()
        # Map relaxed disjunctions -> list of unfixed disjuncts
        root_util_blk.disjunction_to_unfixed_disjuncts = ComponentMap()

        # Preprocess the active disjunctions
        for disjunction in root_util_blk.disjunction_list:
            assert disjunction.active

            disjuncts_fixed_True = []
            disjuncts_fixed_False = []
            unfixed_disjuncts = []

            # categorize the disjuncts in the disjunction
            for disjunct in disjunction.disjuncts:
                if disjunct.indicator_var.fixed:
                    if disjunct.indicator_var.value:
                        disjuncts_fixed_True.append(disjunct)
                    elif disjunct.indicator_var.value == False:
                        disjuncts_fixed_False.append(disjunct)
                else:
                    unfixed_disjuncts.append(disjunct)

            # update disjunct lists for predetermined disjunctions
            if len(disjuncts_fixed_False) == len(disjunction.disjuncts) - 1:
                # all but one disjunct in the disjunction is fixed to False.
                # Remaining one must be true. If not already fixed to True, do
                # so.
                if not disjuncts_fixed_True:
                    disjuncts_fixed_True = unfixed_disjuncts[0]
                    unfixed_disjuncts = []
                    disjuncts_fixed_True.indicator_var.fix(True)
            elif disjuncts_fixed_True and disjunction.xor:
                assert len(disjuncts_fixed_True) == 1, ("XOR (only one True) "
                                                        "violated: %s" %
                                                        disjunction.name)
                disjuncts_fixed_False.extend(unfixed_disjuncts)
                unfixed_disjuncts = []

            # Make sure disjuncts fixed to False are properly deactivated.
            for disjunct in disjuncts_fixed_False:
                disjunct.deactivate()

            # Deactivate nonlinear constraints in unfixed disjuncts
            for disjunct in unfixed_disjuncts:
                nonlinear_constraints_in_disjunct = [
                    constr for constr in disjunct.component_data_objects(
                        Constraint, active=True)
                    if constr.body.polynomial_degree() not in _linear_degrees]
                for constraint in nonlinear_constraints_in_disjunct:
                    constraint.deactivate()
                if nonlinear_constraints_in_disjunct:
                    # TODO might be worthwhile to log number of nonlinear
                    # constraints in each disjunction for later branching
                    # purposes
                    root_util_blk.disjunct_to_nonlinear_constraints[
                        disjunct] = nonlinear_constraints_in_disjunct

            root_util_blk.disjunction_to_unfixed_disjuncts[
                disjunction] = unfixed_disjuncts
            pass

        # Add the BigM suffix if it does not already exist. Used later during
        # nonlinear constraint activation.
        if not hasattr(root_node, 'BigM'):
            root_node.BigM = Suffix()

        # Set up the priority queue
        queue = self.bb_queue = []
        self.created_nodes = 0
        unbranched_disjunction_indices = [
            i for i, disjunction in enumerate(root_util_blk.disjunction_list)
            if disjunction in root_util_blk.disjunction_to_unfixed_disjuncts]
        sort_tuple = BBNodeData(
            obj_lb=float('-inf'),
            obj_ub=float('inf'),
            is_screened=False,
            is_evaluated=False,
            num_unbranched_disjunctions=len(unbranched_disjunction_indices),
            node_count=0,
            unbranched_disjunction_indices=unbranched_disjunction_indices,
        )
        heappush(queue, (sort_tuple, root_node))

        # Do the branch and bound
        while len(queue) > 0:
            # visit the top node on the heap
            node_data, node_model = heappop(queue)
            config.logger.info("Nodes: %s LB %.10g Unbranched %s" % (
                self.explored_nodes, node_data.obj_lb,
                node_data.num_unbranched_disjunctions))

            # Check time limit
            if self.reached_time_limit(config):
                no_feasible_soln = float('inf')
                self.LB = node_data.obj_lb if \
                                 solve_data.objective_sense == minimize else \
                                 -no_feasible_soln
                self.UB = no_feasible_soln if \
                                 solve_data.objective_sense == minimize else \
                                 -node_data.obj_lb
                config.logger.info(
                    'Final bound values: LB: {}  UB: {}'.
                    format(self.LB, self.UB))
                return self._get_final_results_object()

            # Handle current node
            if not node_data.is_screened:
                # Node has not been evaluated.
                self.explored_nodes += 1
                new_node_data = self._prescreen_node(node_data, node_model,
                                                     config)
                # replace with updated node data
                heappush(queue, (new_node_data, node_model))
            elif node_data.obj_lb < node_data.obj_ub - \
                 config.bound_tolerance and not node_data.is_evaluated:
                # Node has not been fully evaluated.
                # Note: infeasible and unbounded nodes will skip this condition,
                # because of strict inequality
                new_node_data = self._evaluate_node(node_data, node_model,
                                                    config)
                # replace with updated node data
                heappush(queue, (new_node_data, node_model))
            elif node_data.num_unbranched_disjunctions == 0 or \
                 node_data.obj_lb == float('inf'):
                # We have reached a leaf node, or the best available node is
                # infeasible.

                # Update the incumbent and put it in the original model
                self.update_incumbent(
                    node_model.component(self.original_util_block.name))
                self._transfer_incumbent_to_original_model(config.logger)

                self.LB = node_data.obj_lb if \
                                 self.objective_sense == minimize else \
                                 -node_data.obj_ub
                self.UB = node_data.obj_ub if \
                                 self.objective_sense == minimize else \
                                 -node_data.obj_lb
                self.iteration = self.explored_nodes
                if node_data.obj_lb == float('inf'):
                    self.pyomo_results.solver.\
                        termination_condition = tc.infeasible
                elif node_data.obj_ub == float('-inf'):
                    self.pyomo_results.solver.\
                        termination_condition = tc.unbounded
                else:
                    self.pyomo_results.solver.\
                        termination_condition = tc.optimal
                return self._get_final_pyomo_results_object()
            else:
                self._branch_on_node(node_data, node_model, config)

    def _branch_on_node(self, node_data, node_model, config):
        node_utils = node_model.component(self.original_util_block.name)

        # Keeping the naive branch selection
        disjunction_to_branch_idx = node_data.unbranched_disjunction_indices[0]
        disjunction_to_branch = node_utils.disjunction_list[
            disjunction_to_branch_idx]
        num_unfixed_disjuncts = len(
            node_utils.disjunction_to_unfixed_disjuncts[disjunction_to_branch])
        config.logger.info("Branching on disjunction %s" %
                           disjunction_to_branch.name)
        node_count = self.created_nodes
        newly_created_nodes = 0

        for disjunct_index_to_fix_True in range(num_unfixed_disjuncts):
            # Create a new branch for each unfixed disjunct
            child_model = node_model.clone()
            child_utils = child_model.component(node_utils.name)
            child_disjunction_to_branch = child_utils.\
                                          disjunction_list[
                                              disjunction_to_branch_idx]
            child_unfixed_disjuncts = child_utils.\
                                      disjunction_to_unfixed_disjuncts[
                                          child_disjunction_to_branch]
            for idx, child_disjunct in enumerate(child_unfixed_disjuncts):
                if idx == disjunct_index_to_fix_True:
                    child_disjunct.indicator_var.fix(True)
                else:
                    child_disjunct.deactivate()
            if not child_disjunction_to_branch.xor:
                raise NotImplementedError("We still need to add support for "
                                          "non-XOR disjunctions.")
            # This requires adding all combinations of activation status among
            # unfixed_disjuncts Reactivate nonlinear constraints in the
            # newly-fixed child disjunct
            fixed_True_disjunct = child_unfixed_disjuncts[
                disjunct_index_to_fix_True]
            for constr in child_utils.\
                disjunct_to_nonlinear_constraints.get(fixed_True_disjunct, ()):
                constr.activate()
                child_model.BigM[constr] = 1  # set arbitrary BigM (ok, because
                                              # we fix corresponding Y=True)

            del child_utils.disjunction_to_unfixed_disjuncts[
                child_disjunction_to_branch]
            for child_disjunct in child_unfixed_disjuncts:
                child_utils.disjunct_to_nonlinear_constraints.pop(
                    child_disjunct, None)

            newly_created_nodes += 1
            child_node_data = node_data._replace(
                is_screened=False,
                is_evaluated=False,
                num_unbranched_disjunctions=node_data.\
                num_unbranched_disjunctions - 1,
                node_count=node_count + newly_created_nodes,
                unbranched_disjunction_indices=node_data.\
                unbranched_disjunction_indices[1:],
                obj_ub=float('inf'),
            )
            heappush(self.bb_queue, (child_node_data, child_model))

        self.created_nodes += newly_created_nodes

        config.logger.info("Added %s new nodes with %s relaxed disjunctions to "
                           "the heap. Size now %s." % (
                               num_unfixed_disjuncts,
                               node_data.num_unbranched_disjunctions - 1,
                               len(self.bb_queue)))

    def _prescreen_node(self, node_data, node_model, config):
        # Check node for satisfiability if sat-solver is enabled
        if config.check_sat and satisfiable(node_model, config.logger) is False:
            if node_data.node_count == 0:
                config.logger.info("Root node is not satisfiable. Problem is "
                                   "infeasible.")
            else:
                config.logger.info("SAT solver pruned node %s" %
                                   node_data.node_count)
            new_lb = new_ub = float('inf')
        else:
            # Solve model subproblem
            if config.solve_local_rnGDP:
                config.logger.debug(
                    "Screening node %s with LB %.10g and %s inactive "
                    "disjunctions." % (node_data.node_count, node_data.obj_lb,
                                       node_data.num_unbranched_disjunctions))
                new_lb, new_ub = self._solve_local_rnGDP_subproblem(node_model,
                                                                    config)
            else:
                new_lb, new_ub = float('-inf'), float('inf')
            new_lb = max(node_data.obj_lb, new_lb)

        new_node_data = node_data._replace(obj_lb=new_lb, obj_ub=new_ub,
                                           is_screened=True)
        return new_node_data

    def _evaluate_node(self, node_data, node_model, config):
        # Solve model subproblem
        config.logger.info(
            "Exploring node %s with LB %.10g UB %.10g and %s inactive "
            "disjunctions." % (node_data.node_count, node_data.obj_lb,
                               node_data.obj_ub,
                               node_data.num_unbranched_disjunctions))
        new_lb, new_ub = self._solve_rnGDP_subproblem(node_model, config)
        new_node_data = node_data._replace(obj_lb=new_lb, obj_ub=new_ub,
                                           is_evaluated=True)
        return new_node_data

    def _solve_rnGDP_subproblem(self, model, config):
        subproblem = TransformationFactory('gdp.bigm').create_using(model)
        obj_sense_correction = self.objective_sense != minimize
        model_utils = model.component(self.original_util_block.name)
        subprob_utils = subproblem.component(
            self.original_util_block.name)

        try:
            with SuppressInfeasibleWarning():
                try:
                    fbbt(subproblem, integer_tol=config.integer_tolerance)
                except InfeasibleConstraintException:
                    # copy variable values, even if errored
                    copy_var_list_values(
                        from_list=subprob_utils.algebraic_variable_list,
                        to_list=model_utils.algebraic_variable_list,
                        config=config, ignore_integrality=True
                    )
                    return float('inf'), float('inf')
                minlp_args = dict(config.minlp_solver_args)
                if (config.time_limit is not None and
                    config.minlp_solver == 'gams'):
                    elapsed = get_main_elapsed_time(self.timing)
                    remaining = max(config.time_limit - elapsed, 1)
                    minlp_args['add_options'] = minlp_args.get('add_options',
                                                               [])
                    minlp_args['add_options'].append('option reslim=%s;' %
                                                     remaining)
                result = SolverFactory(config.minlp_solver).solve(subproblem,
                                                                  **minlp_args)
        except RuntimeError as e:
            config.logger.warning(
                "Solver encountered RuntimeError. Treating as infeasible. "
                "Msg: %s\n%s" % (str(e), traceback.format_exc()))
            copy_var_list_values(  # copy variable values, even if errored
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('inf'), float('inf')

        term_cond = result.solver.termination_condition
        if term_cond == tc.optimal:
            assert result.solver.status is SolverStatus.ok
            lb = result.problem.lower_bound if not obj_sense_correction else \
                 -result.problem.upper_bound
            ub = result.problem.upper_bound if not obj_sense_correction else \
                 -result.problem.lower_bound
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config,
            )
            return lb, ub
        elif term_cond == tc.locallyOptimal or term_cond == tc.feasible:
            assert result.solver.status is SolverStatus.ok
            lb = result.problem.lower_bound if not obj_sense_correction else \
                 -result.problem.upper_bound
            ub = result.problem.upper_bound if not obj_sense_correction else \
                 -result.problem.lower_bound
            # TODO handle LB absent
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config,
            )
            return lb, ub
        elif term_cond == tc.unbounded:
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('-inf'), float('-inf')
        elif term_cond == tc.infeasible:
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('inf'), float('inf')
        else:
            config.logger.warning("Unknown termination condition of %s. "
                                  "Treating as infeasible." % term_cond)
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('inf'), float('inf')


    def _solve_local_rnGDP_subproblem(self, model, config):
        # TODO: The returns of this method should be improved. Currently, it
        # returns trivial bounds (LB, UB) = (-inf, inf) if there is an error in
        # the solve, if the problem is infeasible, or if the problem is
        # unbounded.
        subproblem = TransformationFactory('gdp.bigm').create_using(model)
        obj_sense_correction = self.objective_sense != minimize
        subprob_utils = subproblem.component(
            self.original_util_block.name)
        model_utils = model.component(self.original_util_block.name)

        try:
            with SuppressInfeasibleWarning():
                result = SolverFactory(config.local_minlp_solver).solve(
                    subproblem, **config.local_minlp_solver_args)
        except RuntimeError as e:
            config.logger.warning(
                "Solver encountered RuntimeError. Treating as infeasible. "
                "Msg: %s\n%s" % (str(e), traceback.format_exc()))
            copy_var_list_values(  # copy variable values, even if errored
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('-inf'), float('inf')

        term_cond = result.solver.termination_condition
        if term_cond == tc.optimal:
            assert result.solver.status is SolverStatus.ok
            ub = result.problem.upper_bound if not obj_sense_correction else \
                 -result.problem.lower_bound
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config,
            )
            return float('-inf'), ub
        elif term_cond == tc.locallyOptimal or term_cond == tc.feasible:
            assert result.solver.status is SolverStatus.ok
            ub = result.problem.upper_bound if not obj_sense_correction else \
                 -result.problem.lower_bound
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config,
            )
            return float('-inf'), ub
        elif term_cond == tc.unbounded:
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('-inf'), float('-inf')
        elif term_cond == tc.infeasible:
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            # [ESJ 7/11/22]: Making this float('inf'), float('inf') doesn't
            # break tests, and makes more sense. But I'm not sure why it's not
            # that way already?
            return float('-inf'), float('inf')
        else:
            config.logger.warning("Unknown termination condition of %s. "
                                  "Treating as infeasible." % term_cond)
            copy_var_list_values(
                from_list=subprob_utils.algebraic_variable_list,
                to_list=model_utils.algebraic_variable_list,
                config=config, ignore_integrality=True
            )
            return float('-inf'), float('inf')

GDP_LBB_Solver.solve.__doc__ = add_docstring_list(
    GDP_LBB_Solver.solve.__doc__, GDP_LBB_Solver.CONFIG, indent_by=8)
