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

"""Tests the zero sum propagation module."""

import pyomo.common.unittest as unittest
from pyomo.environ import (
    ConcreteModel,
    Constraint,
    TransformationFactory,
    Var,
    NonNegativeReals,
    NonPositiveReals,
    Binary,
)


class TestZeroSumPropagate(unittest.TestCase):
    """Tests zero sum propagation."""

    def test_non_negative_propagate(self):
        """Test zero sum propagation for non-negative vars."""
        m = ConcreteModel()
        m.v1 = Var(initialize=0)
        m.v2 = Var(initialize=2, domain=NonNegativeReals)
        m.v3 = Var(initialize=3, domain=NonNegativeReals)
        m.v4 = Var(initialize=4, domain=NonNegativeReals)
        m.c1 = Constraint(expr=m.v1 == m.v2 + m.v3 + m.v4)
        m.v1.fix()
        # Because v1 = 0, all the other values must be zero
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertTrue(m.v1.fixed)
        self.assertTrue(m.v2.fixed)
        self.assertTrue(m.v3.fixed)
        self.assertTrue(m.v4.fixed)
        del m

        m = ConcreteModel()
        m.v1 = Var(initialize=0)
        m.v2 = Var(initialize=2, domain=NonNegativeReals)
        m.v3 = Var(initialize=3, domain=NonNegativeReals)
        m.v4 = Var(initialize=4, domain=NonNegativeReals)
        m.c1 = Constraint(expr=m.v2 + m.v3 + m.v4 == m.v1)
        m.v1.fix()
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertTrue(m.v1.fixed)
        self.assertTrue(m.v2.fixed)
        self.assertTrue(m.v3.fixed)
        self.assertTrue(m.v4.fixed)
        del m

        m = ConcreteModel()
        m.v1 = Var(initialize=0)
        m.v2 = Var(initialize=2, domain=NonNegativeReals)
        m.v3 = Var(initialize=3, domain=NonNegativeReals)
        m.v4 = Var(initialize=4, domain=NonNegativeReals)
        m.c1 = Constraint(expr=m.v1 >= m.v2 + m.v3 + m.v4)
        m.v1.fix()
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertTrue(m.v1.fixed)
        self.assertTrue(m.v2.fixed)
        self.assertTrue(m.v3.fixed)
        self.assertTrue(m.v4.fixed)

    def test_not_transform_improperly(self):
        """Tests that invalid constraints are not transformed."""
        m = ConcreteModel()
        m.v1 = Var(initialize=0, domain=Binary)
        m.c1 = Constraint(expr=-1 * m.v1 <= 0)
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertFalse(m.v1.fixed)

    def test_non_positive_propagate(self):
        """Tests zero sum propagation for non-positive vars."""
        m = ConcreteModel()
        m.v1 = Var(initialize=0)
        m.v2 = Var(initialize=-2, domain=NonPositiveReals)
        m.v3 = Var(initialize=-3, domain=NonPositiveReals)
        m.v4 = Var(initialize=-4, domain=NonPositiveReals)
        m.c1 = Constraint(expr=m.v1 == m.v2 + m.v3 + m.v4)
        m.v1.fix()
        # Because v1 = 0, all the other values must be zero
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertTrue(m.v1.fixed)
        self.assertTrue(m.v2.fixed)
        self.assertTrue(m.v3.fixed)
        self.assertTrue(m.v4.fixed)
        del m

        m = ConcreteModel()
        m.v1 = Var(initialize=0)
        m.v2 = Var(initialize=-2, domain=NonPositiveReals)
        m.v3 = Var(initialize=-3, domain=NonPositiveReals)
        m.v4 = Var(initialize=-4, domain=NonPositiveReals)
        m.c1 = Constraint(expr=m.v2 + m.v3 + m.v4 == m.v1)
        m.v1.fix()
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertTrue(m.v1.fixed)
        self.assertTrue(m.v2.fixed)
        self.assertTrue(m.v3.fixed)
        self.assertTrue(m.v4.fixed)
        del m

        m = ConcreteModel()
        m.v1 = Var(initialize=0)
        m.v2 = Var(initialize=-2, domain=NonPositiveReals)
        m.v3 = Var(initialize=-3, domain=NonPositiveReals)
        m.v4 = Var(initialize=-4, domain=NonPositiveReals)
        m.c1 = Constraint(expr=m.v1 <= m.v2 + m.v3 + m.v4)
        m.v1.fix()
        TransformationFactory('contrib.propagate_zero_sum').apply_to(m)
        self.assertTrue(m.v1.fixed)
        self.assertTrue(m.v2.fixed)
        self.assertTrue(m.v3.fixed)
        self.assertTrue(m.v4.fixed)


if __name__ == '__main__':
    unittest.main()
