# ____________________________________________________________________________________
#
# 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 (
    TransformationFactory,
    Var,
    NonNegativeReals,
    Constraint,
    Objective,
    Block,
    value,
)

from pyomo.common.modeling import unique_component_name
from pyomo.core.plugins.transform.hierarchy import NonIsomorphicTransformation
from pyomo.common.config import ConfigBlock, ConfigValue
from pyomo.core.base import ComponentUID
from pyomo.common.deprecation import deprecation_warning


def target_list(x):
    deprecation_msg = (
        "In future releases ComponentUID targets will no "
        "longer be supported in the core.add_slack_variables "
        "transformation. Specify targets as a Constraint or "
        "list of Constraints."
    )
    if isinstance(x, ComponentUID):
        if deprecation_msg:
            deprecation_warning(deprecation_msg, version='5.7.1')
            # only emit the message once
            deprecation_msg = None
        # [ESJ 07/15/2020] We have to just pass it through because we need the
        # instance in order to be able to do anything about it...
        return [x]
    elif getattr(x, 'ctype', None) is Constraint:
        return [x]
    elif hasattr(x, '__iter__'):
        ans = []
        for i in x:
            if isinstance(i, ComponentUID):
                if deprecation_msg:
                    deprecation_warning(deprecation_msg, version='5.7.1')
                    deprecation_msg = None
                # same as above...
                ans.append(i)
            elif getattr(i, 'ctype', None) is Constraint:
                ans.append(i)
            else:
                raise ValueError(
                    "Expected Constraint or list of Constraints."
                    "\n\tReceived %s" % (type(i),)
                )
        return ans
    else:
        raise ValueError(
            "Expected Constraint or list of Constraints.\n\tReceived %s" % (type(x),)
        )


import logging

logger = logging.getLogger('pyomo.core')


@TransformationFactory.register(
    'core.add_slack_variables',
    doc="Create a model where we add slack variables to every constraint "
    "and add new objective penalizing the sum of the slacks",
)
class AddSlackVariables(NonIsomorphicTransformation):
    """
    This plugin adds slack variables to every constraint or to the constraints
    specified in targets.
    """

    CONFIG = ConfigBlock("core.add_slack_variables")
    CONFIG.declare(
        'targets',
        ConfigValue(
            default=None,
            domain=target_list,
            description="target or list of targets to which slacks will be added",
            doc="This specifies the list of Constraints to add slack variables to.",
        ),
    )

    def __init__(self, **kwds):
        kwds['name'] = "add_slack_vars"
        super(AddSlackVariables, self).__init__(**kwds)

    def _apply_to(self, instance, **kwds):
        self._apply_to_impl(instance, **kwds)

    def _apply_to_impl(self, instance, **kwds):
        config = self.CONFIG(kwds.pop('options', {}))
        config.set_value(kwds)
        targets = config.targets

        if targets is None:
            constraintDatas = instance.component_data_objects(
                Constraint, descend_into=True
            )
        else:
            constraintDatas = []
            for t in targets:
                if isinstance(t, ComponentUID):
                    cons = t.find_component(instance)
                    if cons.is_indexed():
                        for i in cons:
                            constraintDatas.append(cons[i])
                    else:
                        constraintDatas.append(cons)
                else:
                    # we know it's a constraint because that's all we let
                    # through the config block validation.
                    if t.is_indexed():
                        for i in t:
                            constraintDatas.append(t[i])
                    else:
                        constraintDatas.append(t)

        # deactivate the objective
        for o in instance.component_data_objects(Objective):
            o.deactivate()

        # create block where we can add slack variables safely
        xblockname = unique_component_name(instance, "_core_add_slack_variables")
        instance.add_component(xblockname, Block())
        xblock = instance.component(xblockname)

        obj_expr = 0
        for cons in constraintDatas:
            if (cons.lower is not None and cons.upper is not None) and value(
                cons.lower
            ) > value(cons.upper):
                # this is a structural infeasibility so slacks aren't going to
                # help:
                raise RuntimeError(
                    "Lower bound exceeds upper bound in constraint %s" % cons.name
                )
            if not cons.active:
                continue
            cons_name = cons.getname(fully_qualified=True)
            lower = cons.lower
            body = cons.body
            upper = cons.upper
            if lower is not None:
                # we add positive slack variable to body:
                # declare positive slack
                varName = "_slack_plus_" + cons_name
                posSlack = Var(within=NonNegativeReals)
                xblock.add_component(varName, posSlack)
                # add positive slack to body expression
                body += posSlack
                # penalize slack in objective
                obj_expr += posSlack
            if upper is not None:
                # we subtract a positive slack variable from the body:
                # declare slack
                varName = "_slack_minus_" + cons_name
                negSlack = Var(within=NonNegativeReals)
                xblock.add_component(varName, negSlack)
                # add negative slack to body expression
                body -= negSlack
                # add slack to objective
                obj_expr += negSlack
            cons.set_value((lower, body, upper))
        # make a new objective that minimizes sum of slack variables
        xblock._slack_objective = Objective(expr=obj_expr)
