# ____________________________________________________________________________________
#
# 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 for applying basic steps."""

import pyomo.common.unittest as unittest
from pyomo.core import Constraint, BooleanVar, SortComponents
from pyomo.gdp.basic_step import apply_basic_step
from pyomo.repn import generate_standard_repn
import pyomo.gdp.tests.models as models
import pyomo.gdp.tests.common_tests as ct
from pyomo.common.fileutils import import_file

from os.path import abspath, dirname, normpath, join

currdir = dirname(abspath(__file__))
exdir = normpath(join(currdir, '..', '..', '..', 'examples', 'gdp'))


class TestBasicStep(unittest.TestCase):
    """Tests disjunctive basic steps."""

    def test_improper_basic_step(self):
        model_builder = import_file(join(exdir, 'two_rxn_lee', 'two_rxn_model.py'))
        m = model_builder.build_model()
        m.basic_step = apply_basic_step([m.reactor_choice, m.max_demand])
        for disj in m.basic_step.disjuncts.values():
            self.assertEqual(disj.improper_constraints[1].body.polynomial_degree(), 2)
            self.assertEqual(disj.improper_constraints[1].lower, None)
            self.assertEqual(disj.improper_constraints[1].upper, 2)
            self.assertEqual(len(disj.improper_constraints), 1)
        self.assertFalse(m.max_demand.active)

    def test_improper_basic_step_linear(self):
        model_builder = import_file(join(exdir, 'two_rxn_lee', 'two_rxn_model.py'))
        m = model_builder.build_model(use_mccormick=True)
        m.basic_step = apply_basic_step(
            [m.reactor_choice, m.max_demand, m.mccormick_1, m.mccormick_2]
        )
        for disj in m.basic_step.disjuncts.values():
            self.assertIs(disj.improper_constraints[1].body, m.P)
            self.assertEqual(disj.improper_constraints[1].lower, None)
            self.assertEqual(disj.improper_constraints[1].upper, 2)
            self.assertEqual(disj.improper_constraints[2].body.polynomial_degree(), 1)
            self.assertEqual(disj.improper_constraints[2].lower, None)
            self.assertEqual(disj.improper_constraints[2].upper, 0)
            self.assertEqual(disj.improper_constraints[3].body.polynomial_degree(), 1)
            self.assertEqual(disj.improper_constraints[3].lower, None)
            self.assertEqual(disj.improper_constraints[3].upper, 0)
            self.assertEqual(len(disj.improper_constraints), 3)
        self.assertFalse(m.max_demand.active)
        self.assertFalse(m.mccormick_1.active)
        self.assertFalse(m.mccormick_2.active)

    def check_constraint_body(self, m, constraint, constant):
        self.assertIsNone(constraint.lower)
        self.assertEqual(constraint.upper, 0)
        repn = generate_standard_repn(constraint.body)
        self.assertEqual(repn.constant, constant)
        self.assertEqual(len(repn.linear_vars), 2)
        ct.check_linear_coef(self, repn, m.a, -1)
        ct.check_linear_coef(self, repn, m.x, 1)

    def check_after_improper_basic_step(self, m):
        for disj in m.basic_step.disjuncts.values():
            self.assertEqual(len(disj.improper_constraints), 1)
            cons = disj.improper_constraints[1]
            self.check_constraint_body(m, cons, -1)

    def test_improper_basic_step_simpleConstraint(self):
        m = models.makeTwoTermDisj()
        m.simple = Constraint(expr=m.x <= m.a + 1)

        m.basic_step = apply_basic_step([m.disjunction, m.simple])
        self.check_after_improper_basic_step(m)

        self.assertFalse(m.simple.active)
        self.assertFalse(m.disjunction.active)

    def test_improper_basic_step_constraintData(self):
        m = models.makeTwoTermDisj()

        @m.Constraint([1, 2])
        def indexed(m, i):
            return m.x <= m.a + i

        m.basic_step = apply_basic_step([m.disjunction, m.indexed[1]])
        self.check_after_improper_basic_step(m)

        self.assertFalse(m.indexed[1].active)
        self.assertTrue(m.indexed[2].active)
        self.assertFalse(m.disjunction.active)

    def test_improper_basic_step_indexedConstraint(self):
        m = models.makeTwoTermDisj()

        @m.Constraint([1, 2])
        def indexed(m, i):
            return m.x <= m.a + i

        m.basic_step = apply_basic_step([m.disjunction, m.indexed])
        for disj in m.basic_step.disjuncts.values():
            self.assertEqual(len(disj.improper_constraints), 2)
            cons = disj.improper_constraints[1]
            self.check_constraint_body(m, cons, -1)

            cons = disj.improper_constraints[2]
            self.check_constraint_body(m, cons, -2)

    def test_indicator_var_references(self):
        m = models.makeTwoTermDisj()
        m.simple = Constraint(expr=m.x <= m.a + 1)

        m.basic_step = apply_basic_step([m.disjunction, m.simple])

        refs = [
            v
            for v in m.basic_step.component_data_objects(
                BooleanVar, sort=SortComponents.deterministic
            )
        ]
        self.assertEqual(len(refs), 2)
        self.assertIs(refs[0][None], m.d[0].indicator_var)
        self.assertIs(refs[1][None], m.d[1].indicator_var)

    def test_arg_errors(self):
        m = models.makeTwoTermDisj()
        m.simple = Constraint(expr=m.x <= m.a + 1)

        with self.assertRaisesRegex(
            ValueError,
            'apply_basic_step only accepts a list containing '
            'Disjunctions or Constraints',
        ):
            apply_basic_step([m.disjunction, m.simple, m.x])
        with self.assertRaisesRegex(
            ValueError,
            'apply_basic_step: argument list must contain at least one Disjunction',
        ):
            apply_basic_step([m.simple, m.simple])


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