# ____________________________________________________________________________________
#
# 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
import pyomo.environ as pyo
from pyomo.contrib.mpc.data.get_cuid import get_indexed_cuid


class TestGetCUID(unittest.TestCase):
    def _make_model(self):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=[0.0, 0.1, 0.2])
        m.space = pyo.Set(initialize=[1.0, 1.5, 2.0])
        m.comp = pyo.Set(initialize=["A", "B"])
        m.var = pyo.Var(m.time, m.comp, initialize=1.0)
        m.txc_var = pyo.Var(m.time, m.space, m.comp, initialize=2.0)

        @m.Block(m.time, m.space)
        def b(b, t, x):
            b.bvar1 = pyo.Var(initialize=3.0)
            b.bvar2 = pyo.Var(m.comp, initialize=3.0)

        return m

    def test_get_cuid(self):
        m = self._make_model()

        pred_cuid = pyo.ComponentUID(m.var[:, "A"])
        self.assertEqual(get_indexed_cuid(m.var[:, "A"]), pred_cuid)
        self.assertEqual(get_indexed_cuid(pyo.Reference(m.var[:, "A"])), pred_cuid)
        self.assertEqual(get_indexed_cuid("var[*,A]"), pred_cuid)
        self.assertEqual(get_indexed_cuid("var[*,'A']"), pred_cuid)
        self.assertEqual(get_indexed_cuid(m.var[0, "A"], sets=(m.time,)), pred_cuid)

    def test_get_cuid_twosets(self):
        m = self._make_model()

        pred_cuid = pyo.ComponentUID(m.b[:, :].bvar2["A"])
        self.assertEqual(get_indexed_cuid(m.b[:, :].bvar2["A"]), pred_cuid)
        self.assertEqual(
            get_indexed_cuid(pyo.Reference(m.b[:, :].bvar2["A"])), pred_cuid
        )
        self.assertEqual(get_indexed_cuid("b[*,*].bvar2[A]"), pred_cuid)
        self.assertEqual(
            get_indexed_cuid(m.b[0, 1].bvar2["A"], sets=(m.time, m.space)), pred_cuid
        )

    def test_get_cuid_dereference(self):
        m = self._make_model()
        m.ref = pyo.Reference(m.var[:, "A"])
        m.ref2 = pyo.Reference(m.ref)

        pred_cuid = pyo.ComponentUID(m.var[:, "A"])

        # ref is attached to the model, so by default we don't "dereference"
        self.assertNotEqual(get_indexed_cuid(m.ref), pred_cuid)
        self.assertEqual(get_indexed_cuid(m.ref), pyo.ComponentUID(m.ref[:]))

        # If we use dereference=True, we do dereference and get the CUID of
        # the underlying slice.
        self.assertEqual(get_indexed_cuid(m.ref, dereference=True), pred_cuid)

        # However, we only dereference once, so a reference-to-reference
        # does not reveal the underlying slice (of the original reference)
        self.assertNotEqual(get_indexed_cuid(m.ref2, dereference=True), pred_cuid)
        self.assertEqual(
            get_indexed_cuid(m.ref2, dereference=True), pyo.ComponentUID(m.ref[:])
        )

        # But if we use dereference=2, we allow two dereferences, and get
        # the original slice
        self.assertEqual(get_indexed_cuid(m.ref2, dereference=2), pred_cuid)

    def test_get_cuid_context(self):
        m = self._make_model()
        top = pyo.ConcreteModel()
        top.m = m

        pred_cuid = pyo.ComponentUID(m.var[:, "A"], context=m)
        self.assertEqual(get_indexed_cuid(m.var[:, "A"], context=m), pred_cuid)
        self.assertEqual(
            get_indexed_cuid(pyo.Reference(m.var[:, "A"]), context=m), pred_cuid
        )

        # This is what we would expect without a context arg
        full_cuid = pyo.ComponentUID(m.var[:, "A"])
        self.assertNotEqual(get_indexed_cuid("m.var[*,A]"), pred_cuid)
        self.assertEqual(get_indexed_cuid("m.var[*,A]"), full_cuid)

        msg = "Context is not allowed"
        with self.assertRaisesRegex(ValueError, msg):
            # Passing context with a string raises a reasonable
            # error from the CUID constructor
            get_indexed_cuid("m.var[*,A]", context=m)

        self.assertEqual(
            get_indexed_cuid(m.var[0, "A"], sets=(m.time,), context=m), pred_cuid
        )


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