# ____________________________________________________________________________________
#
# 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 pyomo.common.unittest as unittest
from pyomo.contrib.cp.interval_var import (
    IntervalVar,
    IntervalVarTimePoint,
    IntervalVarLength,
    IntervalVarPresence,
)
from pyomo.core.expr import GetItemExpression, GetAttrExpression
from pyomo.environ import ConcreteModel, Integers, Reference, Set, value, Var


class TestScalarIntervalVar(unittest.TestCase):
    def test_initialize_with_no_data(self):
        m = ConcreteModel()
        m.i = IntervalVar()

        self.assertIsInstance(m.i.start_time, IntervalVarTimePoint)
        self.assertEqual(m.i.start_time.domain, Integers)
        self.assertIsNone(m.i.start_time.lower)
        self.assertIsNone(m.i.start_time.upper)

        self.assertIsInstance(m.i.end_time, IntervalVarTimePoint)
        self.assertEqual(m.i.end_time.domain, Integers)
        self.assertIsNone(m.i.end_time.lower)
        self.assertIsNone(m.i.end_time.upper)

        self.assertIsInstance(m.i.length, IntervalVarLength)
        self.assertEqual(m.i.length.domain, Integers)
        self.assertIsNone(m.i.length.lower)
        self.assertIsNone(m.i.length.upper)

        self.assertIsInstance(m.i.is_present, IntervalVarPresence)

    def test_add_components_that_do_not_belong(self):
        m = ConcreteModel()
        m.i = IntervalVar()

        with self.assertRaisesRegex(
            ValueError,
            "Attempting to declare a block component using the name of a "
            "reserved attribute:\n\tnew_thing",
        ):
            m.i.new_thing = IntervalVar()

    def test_start_and_end_bounds(self):
        m = ConcreteModel()
        m.i = IntervalVar(start=(0, 5))
        self.assertEqual(m.i.start_time.lower, 0)
        self.assertEqual(m.i.start_time.upper, 5)

        m.i.end_time.bounds = (12, 14)

        self.assertEqual(m.i.end_time.lower, 12)
        self.assertEqual(m.i.end_time.upper, 14)

    def test_constant_length_and_start(self):
        m = ConcreteModel()
        m.i = IntervalVar(length=7, start=3)

        self.assertEqual(m.i.length.lower, 7)
        self.assertEqual(m.i.length.upper, 7)

        self.assertEqual(m.i.start_time.lower, 3)
        self.assertEqual(m.i.start_time.upper, 3)

    def test_non_optional(self):
        m = ConcreteModel()
        m.i = IntervalVar(length=2, end=(4, 9), optional=False)

        self.assertEqual(value(m.i.is_present), True)
        self.assertTrue(m.i.is_present.fixed)
        self.assertFalse(m.i.optional)

        # Should also be true by default
        m.i2 = IntervalVar()

        self.assertEqual(value(m.i2.is_present), True)
        self.assertTrue(m.i.is_present.fixed)
        self.assertFalse(m.i2.optional)

    def test_optional(self):
        m = ConcreteModel()
        m.i = IntervalVar(optional=True)

        self.assertFalse(m.i.is_present.fixed)
        self.assertTrue(m.i.optional)

        # Now set to False
        m.i.optional = False
        self.assertEqual(value(m.i.is_present), True)
        self.assertTrue(m.i.is_present.fixed)
        self.assertFalse(m.i.optional)

    def test_is_present_fixed_False(self):
        m = ConcreteModel()
        m.i = IntervalVar(optional=True)

        m.i.is_present.fix(False)
        self.assertTrue(m.i.optional)


class TestIndexedIntervalVar(unittest.TestCase):
    def test_initialize_with_no_data(self):
        m = ConcreteModel()

        m.i = IntervalVar([1, 2])

        for j in [1, 2]:
            self.assertIsInstance(m.i[j].start_time, IntervalVarTimePoint)
            self.assertEqual(m.i[j].start_time.domain, Integers)
            self.assertIsNone(m.i[j].start_time.lower)
            self.assertIsNone(m.i[j].start_time.upper)

            self.assertIsInstance(m.i[j].end_time, IntervalVarTimePoint)
            self.assertEqual(m.i[j].end_time.domain, Integers)
            self.assertIsNone(m.i[j].end_time.lower)
            self.assertIsNone(m.i[j].end_time.upper)

            self.assertIsInstance(m.i[j].length, IntervalVarLength)
            self.assertEqual(m.i[j].length.domain, Integers)
            self.assertIsNone(m.i[j].length.lower)
            self.assertIsNone(m.i[j].length.upper)

            self.assertIsInstance(m.i[j].is_present, IntervalVarPresence)

    def test_constant_length(self):
        m = ConcreteModel()
        m.i = IntervalVar(['a', 'b'], length=45)

        for j in ['a', 'b']:
            self.assertEqual(m.i[j].length.lower, 45)
            self.assertEqual(m.i[j].length.upper, 45)

    def test_rule_based_start(self):
        m = ConcreteModel()

        def start_rule(m, i):
            return (1 - i, 13 + i)

        m.act = IntervalVar([1, 2, 3], start=start_rule, length=4)

        for i in [1, 2, 3]:
            self.assertEqual(m.act[i].start_time.lower, 1 - i)
            self.assertEqual(m.act[i].start_time.upper, 13 + i)

            self.assertEqual(m.act[i].length.lower, 4)
            self.assertEqual(m.act[i].length.upper, 4)

            self.assertFalse(m.act[i].optional)
            self.assertTrue(m.act[i].is_present.fixed)
            self.assertEqual(value(m.act[i].is_present), True)

    def test_optional(self):
        m = ConcreteModel()
        m.act = IntervalVar([1, 2], end=[0, 10], optional=True)

        for i in [1, 2]:
            self.assertTrue(m.act[i].optional)
            self.assertFalse(m.act[i].is_present.fixed)

            self.assertEqual(m.act[i].end_time.lower, 0)
            self.assertEqual(m.act[i].end_time.upper, 10)

        # None doesn't make sense for this:
        with self.assertRaisesRegex(
            ValueError, "Cannot set 'optional' to None: Must be True or False."
        ):
            m.act[1].optional = None

        # We can change it, and that has the correct effect on is_present
        m.act[1].optional = False
        self.assertFalse(m.act[1].optional)
        self.assertTrue(m.act[1].is_present.fixed)

        m.act[1].optional = True
        self.assertTrue(m.act[1].optional)
        self.assertFalse(m.act[1].is_present.fixed)

    def test_optional_rule(self):
        m = ConcreteModel()
        m.idx = Set(initialize=[(4, 2), (5, 2)], dimen=2)

        def optional_rule(m, i, j):
            return i % j == 0

        m.act = IntervalVar(m.idx, optional=optional_rule)
        self.assertTrue(m.act[4, 2].optional)
        self.assertFalse(m.act[5, 2].optional)

    def test_index_by_expr(self):
        m = ConcreteModel()
        m.act = IntervalVar([(1, 2), (2, 1), (2, 2)])
        m.i = Var(domain=Integers)
        m.i2 = Var([1, 2], domain=Integers)

        thing1 = m.act[m.i, 2]
        self.assertIsInstance(thing1, GetItemExpression)
        self.assertEqual(len(thing1.args), 3)
        self.assertIs(thing1.args[0], m.act)
        self.assertIs(thing1.args[1], m.i)
        self.assertEqual(thing1.args[2], 2)

        thing2 = thing1.start_time
        self.assertIsInstance(thing2, GetAttrExpression)
        self.assertEqual(len(thing2.args), 2)
        self.assertIs(thing2.args[0], thing1)
        self.assertEqual(thing2.args[1], 'start_time')

        expr1 = m.act[m.i, 2].start_time.before(m.act[m.i**2, 1].end_time)

    def test_reference(self):
        m = ConcreteModel()
        m.act = IntervalVar([1, 2], end=[0, 10], optional=True)

        thing = Reference(m.act[:].is_present)
        self.assertIs(thing[1], m.act[1].is_present)
        self.assertIs(thing[2], m.act[2].is_present)

        thing = Reference(m.act[:].start_time)
        self.assertIs(thing[1], m.act[1].start_time)
        self.assertIs(thing[2], m.act[2].start_time)

        thing = Reference(m.act[:].end_time)
        self.assertIs(thing[1], m.act[1].end_time)
        self.assertIs(thing[2], m.act[2].end_time)

        thing = Reference(m.act[:].length)
        self.assertIs(thing[1], m.act[1].length)
        self.assertIs(thing[2], m.act[2].length)
