# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 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 math import fabs
from pyomo.contrib.gdpopt.solve_subproblem import solve_subproblem
from pyomo.contrib.gdpopt.util import fix_discrete_problem_solution_in_subproblem
from pyomo.core import value
from pyomo.opt import TerminationCondition as tc


class _OAAlgorithmMixIn:
    def _fix_discrete_soln_solve_subproblem_and_add_cuts(
        self, discrete_prob_util_block, subprob_util_block, config
    ):
        with fix_discrete_problem_solution_in_subproblem(
            discrete_prob_util_block, subprob_util_block, self, config
        ):
            nlp_termination = solve_subproblem(subprob_util_block, self, config)
            if nlp_termination in {tc.optimal, tc.feasible}:
                primal_improved = self._update_bounds_after_solve(
                    'subproblem',
                    primal=value(subprob_util_block.obj.expr),
                    logger=config.logger,
                )
                if primal_improved:
                    self.update_incumbent(subprob_util_block)
                self._add_cuts_to_discrete_problem(
                    subprob_util_block,
                    discrete_prob_util_block,
                    self.objective_sense,
                    config,
                    self.timing,
                )
            elif nlp_termination == tc.unbounded:
                # the whole problem is unbounded, we can stop
                self._update_primal_bound_to_unbounded(config)

        return nlp_termination not in {tc.infeasible, tc.unbounded}

    # Utility used in cut_generation: We saved a map of Disjuncts to the active
    # constraints they contain on the discrete problem util_block, and use it
    # here to find the active constraints under the current discrete
    # solution. Note that this preprocesses not just to be efficient, but
    # because everything on the Disjuncts is deactivated at this point, since
    # we've already transformed the discrete problem to a MILP
    def _get_active_untransformed_constraints(self, util_block, config):
        """Yield constraints in disjuncts where the indicator value is set or
        fixed to True."""
        model = util_block.parent_block()
        # Get active global constraints
        for constr in util_block.global_constraint_list:
            yield constr
        # get all the disjuncts in the original model. Check which ones are
        # True.
        for disj, constr_list in util_block.constraints_by_disjunct.items():
            if fabs(disj.binary_indicator_var.value - 1) <= config.integer_tolerance:
                for constr in constr_list:
                    yield constr
