# ____________________________________________________________________________________
#
# 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 pyomo.core import (
    SortComponents,
    Constraint,
    Objective,
    LogicalConstraint,
    Expression,
)
from pyomo.core.base import TransformationFactory, Suffix, ConstraintList, Integers
from pyomo.core.base.block import Block, TraversalStrategy
from pyomo.common.collections import ComponentMap, ComponentSet
from pyomo.common.modeling import unique_component_name
from pyomo.contrib.gdpopt.discrete_problem_initialize import valid_init_strategies
from pyomo.contrib.gdpopt.util import (
    get_main_elapsed_time,
    move_nonlinear_objective_to_constraints,
)
from pyomo.gdp.disjunct import Disjunct, Disjunction
from pyomo.util.vars_from_expressions import get_vars_from_components


def _get_discrete_problem_and_subproblem(solver, config):
    util_block = solver.original_util_block
    original_model = util_block.parent_block()
    if config.force_subproblem_nlp:
        # We'll need to fix these too
        add_discrete_variable_list(util_block)
    original_obj = move_nonlinear_objective_to_constraints(util_block, config.logger)
    solver.original_obj = original_obj

    # create model to hold the subproblems: We create this first because
    # certain initialization strategies for the discrete problem need it.
    subproblem, subproblem_util_block = get_subproblem(original_model, util_block)

    # create discrete problem--the MILP relaxation
    start = get_main_elapsed_time(solver.timing)
    discrete_problem_util_block = initialize_discrete_problem(
        util_block, subproblem_util_block, config, solver
    )

    config.logger.info(
        'Finished discrete problem initialization in {:.2f}s '
        'and {} iterations \n'.format(
            get_main_elapsed_time(solver.timing) - start,
            solver.initialization_iteration,
        )
    )

    return (discrete_problem_util_block, subproblem_util_block)


def initialize_discrete_problem(util_block, subprob_util_block, config, solver):
    """
    Calls the specified transformation (by default bigm) on the original
    model and removes nonlinear constraints to create a MILP discrete problem.
    """
    config.logger.info("---Starting GDPopt initialization---")
    # clone the original model
    discrete = util_block.parent_block().clone()
    discrete.name = discrete.name + ": discrete problem"

    discrete_problem_util_block = discrete.component(util_block.local_name)
    discrete_problem_util_block.no_good_cuts = ConstraintList()
    discrete_problem_util_block.no_good_disjunctions = Disjunction(Integers)

    # deactivate nonlinear constraints
    for c in discrete.component_data_objects(
        Constraint, active=True, descend_into=(Block, Disjunct)
    ):
        if c.body.polynomial_degree() not in (1, 0):
            c.deactivate()

    # Transform to a MILP
    TransformationFactory(config.discrete_problem_transformation).apply_to(discrete)
    add_transformed_boolean_variable_list(discrete_problem_util_block)
    add_algebraic_variable_list(discrete_problem_util_block, name='all_mip_variables')

    # Call the specified initialization strategy. (We've already validated the
    # input in the config logic, so we know this is okay.)
    init_algorithm = valid_init_strategies.get(config.init_algorithm)
    init_algorithm(
        util_block, discrete_problem_util_block, subprob_util_block, config, solver
    )

    return discrete_problem_util_block


def add_util_block(discrete):
    # create a block to store the cuts
    name = unique_component_name(discrete, '_gdpopt_cuts')
    block = Block()
    discrete.add_component(name, block)

    return block


def add_disjunct_list(util_block):
    model = util_block.parent_block()
    util_block.disjunct_list = list(
        model.component_data_objects(
            ctype=Disjunct,
            active=True,
            descend_into=(Block, Disjunct),
            sort=SortComponents.deterministic,
        )
    )


def add_disjunction_list(util_block):
    model = util_block.parent_block()
    util_block.disjunction_list = list(
        model.component_data_objects(
            ctype=Disjunction,
            active=True,
            descend_into=(Block, Disjunct),
            sort=SortComponents.deterministic,
        )
    )


def add_constraint_list(util_block):
    model = util_block.parent_block()
    util_block.constraint_list = list(
        model.component_data_objects(
            ctype=Constraint,
            active=True,
            descend_into=(Block, Disjunct),
            sort=SortComponents.deterministic,
        )
    )


def add_global_constraint_list(util_block):
    model = util_block.parent_block()
    util_block.global_constraint_list = list(
        model.component_data_objects(
            ctype=Constraint,
            active=True,
            descend_into=Block,
            sort=SortComponents.deterministic,
        )
    )


def add_constraints_by_disjunct(util_block):
    constraints_by_disjunct = util_block.constraints_by_disjunct = {}
    for disj in util_block.disjunct_list:
        cons_list = constraints_by_disjunct[disj] = []
        for cons in disj.component_data_objects(
            Constraint,
            active=True,
            descend_into=Block,
            sort=SortComponents.deterministic,
        ):
            cons_list.append(cons)


def add_algebraic_variable_list(util_block, name=None):
    """
    This collects variables from active Constraints and Objectives. It descends
    into Disjuncts, but does not collect any indicator variables that do not
    appear in algebraic constraints pre-transformation.
    """
    model = util_block.parent_block()
    if name is None:
        name = "algebraic_variable_list"
    setattr(
        util_block,
        name,
        list(
            get_vars_from_components(
                model,
                ctype=(Constraint, Objective),
                descend_into=(Block, Disjunct),
                active=True,
                sort=SortComponents.deterministic,
            )
        ),
    )


def add_discrete_variable_list(util_block):
    lst = util_block.discrete_variable_list = []
    for v in util_block.algebraic_variable_list:
        if v.is_integer():
            lst.append(v)


# Must be collected after list of Disjuncts
def add_boolean_variable_lists(util_block):
    util_block.boolean_variable_list = []
    util_block.non_indicator_boolean_variable_list = []
    for disjunct in util_block.disjunct_list:
        util_block.boolean_variable_list.append(disjunct.indicator_var)
    ind_var_set = ComponentSet(util_block.boolean_variable_list)
    # This will not necessarily include the indicator_vars if it is called
    # before the GDP is transformed to a MIP.
    for v in get_vars_from_components(
        util_block.parent_block(),
        ctype=LogicalConstraint,
        descend_into=(Block, Disjunct),
        active=True,
        sort=SortComponents.deterministic,
    ):
        if v not in ind_var_set:
            util_block.boolean_variable_list.append(v)
            util_block.non_indicator_boolean_variable_list.append(v)


# For the discrete problem, we want the corresponding binaries for all of the
# BooleanVars. This must be called after logical_to_disjunctive has been called.
def add_transformed_boolean_variable_list(util_block):
    util_block.transformed_boolean_variable_list = [
        v.get_associated_binary() for v in util_block.boolean_variable_list
    ]


def get_subproblem(original_model, util_block):
    """Clone the original, and reclassify all the Disjuncts to Blocks.
    We'll also call logical_to_disjunctive and bigm the disjunctive parts in
    case any of the indicator_vars are used in logical constraints and to make
    sure that the rest of the model is algebraic (assuming it was a proper
    GDP to begin with).
    """
    subproblem = original_model.clone()
    subproblem.name = subproblem.name + ": subproblem"

    # Set up dual value reporting
    if not hasattr(subproblem, 'dual'):
        subproblem.dual = Suffix(direction=Suffix.IMPORT)
    elif not isinstance(subproblem.dual, Suffix):
        raise ValueError(
            "The model contains a component called 'dual' that "
            "is not a Suffix. It is of type %s. Please rename "
            "this component, as GDPopt needs dual information to "
            "create cuts." % type(subproblem.dual)
        )
    subproblem.dual.activate()

    # reclassify all the Disjuncts as Blocks and deactivate the Disjunctions. We
    # don't need to add the xor constraints because we're not going to pass
    # infeasible integer solutions to this model.
    for disjunction in subproblem.component_data_objects(
        Disjunction,
        descend_into=(Block, Disjunct),
        descent_order=TraversalStrategy.PostfixDFS,
    ):
        for disjunct in disjunction.disjuncts:
            if disjunct.indicator_var.fixed:
                if not disjunct.indicator_var.value:
                    disjunct.deactivate()
            disjunct.parent_block().reclassify_component_type(disjunct, Block)

        disjunction.deactivate()

    TransformationFactory('contrib.logical_to_disjunctive').apply_to(subproblem)
    # transform any of the Disjuncts we created above with bigm.
    TransformationFactory('gdp.bigm').apply_to(subproblem)

    subproblem_util_block = subproblem.component(util_block.local_name)
    save_initial_values(subproblem_util_block)
    add_transformed_boolean_variable_list(subproblem_util_block)
    subproblem_obj = next(
        subproblem.component_data_objects(Objective, active=True, descend_into=True)
    )
    subproblem_util_block.obj = Expression(expr=subproblem_obj.expr)

    return subproblem, subproblem_util_block


def save_initial_values(subproblem_util_block):
    initial_values = subproblem_util_block.initial_var_values = ComponentMap()
    for v in subproblem_util_block.algebraic_variable_list:
        initial_values[v] = v.value
