# ____________________________________________________________________________________
#
# 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 import (
    TransformationFactory,
    Transformation,
    Block,
    VarList,
    Set,
    SortComponents,
    Objective,
    Constraint,
)
from pyomo.gdp import Disjunct, Disjunction
from pyomo.repn import generate_standard_repn

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


@TransformationFactory.register(
    'gdp.bilinear',
    doc="Creates a disjunctive model where bilinear terms are replaced with disjunctive expressions.",
)
class Bilinear_Transformation(Transformation):
    def __init__(self):
        super(Bilinear_Transformation, self).__init__()

    def _apply_to(self, instance, **kwds):
        # TODO: This data should be stored differently.  We cannot nest this transformation with itself
        if getattr(instance, 'bilinear_data_', None) is None:
            instance.bilinear_data_ = Block()
            instance.bilinear_data_.cache = {}
            instance.bilinear_data_.vlist = VarList()
            instance.bilinear_data_.vlist_boolean = []
            instance.bilinear_data_.IDX = Set()
            instance.bilinear_data_.disjuncts_ = Disjunct(
                instance.bilinear_data_.IDX * [0, 1]
            )
            instance.bilinear_data_.disjunction_data = {}
            instance.bilinear_data_.o_expr = {}
            instance.bilinear_data_.c_body = {}
        #
        # Iterate over all blocks
        #
        for block in instance.block_data_objects(
            active=True, sort=SortComponents.deterministic
        ):
            self._transformBlock(block, instance)

        #
        # WEH: I wish I had a DisjunctList and DisjunctionList object...
        #
        def rule(block, i):
            return instance.bilinear_data_.disjunction_data[i]

        instance.bilinear_data_.disjunction_ = Disjunction(
            instance.bilinear_data_.IDX, rule=rule
        )

    def _transformBlock(self, block, instance):
        for component in block.component_objects(
            Objective, active=True, descend_into=False
        ):
            expr = self._transformExpression(component.expr, instance)
            instance.bilinear_data_.o_expr[id(component)] = component.expr
            component.expr = expr
        for component in block.component_data_objects(
            Constraint, active=True, descend_into=False
        ):
            lb, body, ub = component.to_bounded_expression()
            expr = self._transformExpression(body, instance)
            instance.bilinear_data_.c_body[id(component)] = body
            component.set_value((lb, expr, ub))

    def _transformExpression(self, expr, instance):
        if expr.polynomial_degree() > 2:
            raise ValueError("Cannot transform polynomial terms with degree > 2")
        if expr.polynomial_degree() < 2:
            return expr
        #
        expr = self._replace_bilinear(expr, instance)
        return expr

    def _replace_bilinear(self, expr, instance):
        idMap = {}
        terms = generate_standard_repn(expr, idMap=idMap)
        # Constant
        e = terms.constant
        # Linear terms
        for var, coef in zip(terms.linear_vars, terms.linear_coefs):
            e += coef * var
        # Quadratic terms
        if len(terms.quadratic_coefs) > 0:
            for vars_, coef_ in zip(terms.quadratic_vars, terms.quadratic_coefs):
                #
                if vars_[0].is_binary():
                    v = instance.bilinear_data_.cache.get(
                        (id(vars_[0]), id(vars_[1])), None
                    )
                    if v is None:
                        instance.bilinear_data_.vlist_boolean.append(vars_[0])
                        v = instance.bilinear_data_.vlist.add()
                        instance.bilinear_data_.cache[id(vars_[0]), id(vars_[1])] = v
                        bounds = vars_[1].bounds
                        v.setlb(bounds[0])
                        v.setub(bounds[1])
                        id_ = len(instance.bilinear_data_.vlist)
                        instance.bilinear_data_.IDX.add(id_)
                        # First disjunct
                        d0 = instance.bilinear_data_.disjuncts_[id_, 0]
                        d0.c1 = Constraint(expr=vars_[0] == 1)
                        d0.c2 = Constraint(expr=v == vars_[1])
                        # Second disjunct
                        d1 = instance.bilinear_data_.disjuncts_[id_, 1]
                        d1.c1 = Constraint(expr=vars_[0] == 0)
                        d1.c2 = Constraint(expr=v == 0)
                        # Disjunction
                        instance.bilinear_data_.disjunction_data[id_] = [
                            instance.bilinear_data_.disjuncts_[id_, 0],
                            instance.bilinear_data_.disjuncts_[id_, 1],
                        ]
                        instance.bilinear_data_.disjunction_data[id_] = [
                            instance.bilinear_data_.disjuncts_[id_, 0],
                            instance.bilinear_data_.disjuncts_[id_, 1],
                        ]
                    # The disjunctive variable is the expression
                    e += coef_ * v
                #
                elif vars_[1].is_binary():
                    v = instance.bilinear_data_.cache.get(
                        (id(vars_[1]), id(vars_[0])), None
                    )
                    if v is None:
                        instance.bilinear_data_.vlist_boolean.append(vars_[1])
                        v = instance.bilinear_data_.vlist.add()
                        instance.bilinear_data_.cache[id(vars_[1]), id(vars_[0])] = v
                        bounds = vars_[0].bounds
                        v.setlb(bounds[0])
                        v.setub(bounds[1])
                        id_ = len(instance.bilinear_data_.vlist)
                        instance.bilinear_data_.IDX.add(id_)
                        # First disjunct
                        d0 = instance.bilinear_data_.disjuncts_[id_, 0]
                        d0.c1 = Constraint(expr=vars_[1] == 1)
                        d0.c2 = Constraint(expr=v == vars_[0])
                        # Second disjunct
                        d1 = instance.bilinear_data_.disjuncts_[id_, 1]
                        d1.c1 = Constraint(expr=vars_[1] == 0)
                        d1.c2 = Constraint(expr=v == 0)
                        # Disjunction
                        instance.bilinear_data_.disjunction_data[id_] = [
                            instance.bilinear_data_.disjuncts_[id_, 0],
                            instance.bilinear_data_.disjuncts_[id_, 1],
                        ]
                    # The disjunctive variable is the expression
                    e += coef_ * v
                else:
                    # If neither variable is boolean, just reinsert the original bilinear term
                    e += coef_ * vars_[0] * vars_[1]
        #
        return e
