# ____________________________________________________________________________________
#
# 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.
# ____________________________________________________________________________________

import logging

from pyomo.core.expr import inequality
from pyomo.core.base import (
    Transformation,
    TransformationFactory,
    Constraint,
    Block,
    SortComponents,
    ComponentUID,
)
from pyomo.mpec.complementarity import Complementarity
from pyomo.gdp.disjunct import Disjunct, Disjunction

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


@TransformationFactory.register(
    'mpec.simple_disjunction',
    doc="Disjunctive transformations of complementarity conditions when all variables are non-negative",
)
class MPEC2_Transformation(Transformation):
    def __init__(self):
        super(MPEC2_Transformation, self).__init__()

    def _apply_to(self, instance, **kwds):
        #
        # Setup transformation data
        #
        tdata = instance._transformation_data['mpec.simple_disjunction']
        tdata.compl_cuids = []
        #
        # Iterate over the model finding Complementarity components
        #
        for complementarity in instance.component_objects(
            Complementarity,
            active=True,
            descend_into=(Block, Disjunct),
            sort=SortComponents.deterministic,
        ):
            block = complementarity.parent_block()

            for index in sorted(complementarity.keys()):
                _data = complementarity[index]
                if not _data.active:
                    continue
                #
                _e1 = _data._canonical_expression(_data._args[0])
                _e2 = _data._canonical_expression(_data._args[1])
                if (
                    len(_e1) == 3
                    and len(_e2) == 3
                    and (_e1[0] is None)
                    + (_e1[2] is None)
                    + (_e2[0] is None)
                    + (_e2[2] is None)
                    != 2
                ):
                    raise RuntimeError(
                        "Complementarity condition %s must have exactly two finite bounds"
                        % _data.name
                    )
                if len(_e1) == 3 and _e1[0] is None and _e1[2] is None:
                    #
                    # Swap _e1 and _e2.  The ensures that
                    # only e2 will be an unconstrained expression
                    #
                    _e1, _e2 = _e2, _e1
                if _e2[0] is None and _e2[2] is None:
                    if len(_e1) == 2:
                        _data.c = Constraint(expr=_e1)
                    else:
                        _data.expr1 = Disjunct()
                        _data.expr1.c0 = Constraint(expr=_e1[0] == _e1[1])
                        _data.expr1.c1 = Constraint(expr=_e2[1] >= 0)
                        #
                        _data.expr2 = Disjunct()
                        _data.expr2.c0 = Constraint(expr=_e1[1] == _e1[2])
                        _data.expr2.c1 = Constraint(expr=_e2[1] <= 0)
                        #
                        _data.expr3 = Disjunct()
                        _data.expr3.c0 = Constraint(
                            expr=inequality(_e1[0], _e1[1], _e1[2])
                        )
                        _data.expr3.c1 = Constraint(expr=_e2[1] == 0)
                        _data.complements = Disjunction(
                            expr=(_data.expr1, _data.expr2, _data.expr3)
                        )
                else:
                    if _e1[0] is None:
                        tmp1 = _e1[2] - _e1[1]
                    else:
                        tmp1 = _e1[1] - _e1[0]
                    if _e2[0] is None:
                        tmp2 = _e2[2] - _e2[1]
                    else:
                        tmp2 = _e2[1] - _e2[0]
                    _data.expr1 = Disjunct()
                    _data.expr1.c0 = Constraint(expr=tmp1 >= 0)
                    _data.expr1.c1 = Constraint(expr=tmp2 == 0)
                    #
                    _data.expr2 = Disjunct()
                    _data.expr2.c0 = Constraint(expr=tmp1 == 0)
                    _data.expr2.c1 = Constraint(expr=tmp2 >= 0)
                    #
                    _data.complements = Disjunction(expr=(_data.expr1, _data.expr2))
            tdata.compl_cuids.append(ComponentUID(complementarity))
            block.reclassify_component_type(complementarity, Block)
