# ____________________________________________________________________________________
#
# 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 collections.abc
import pickle

import pyomo.common.unittest as unittest
from pyomo.common.collections import ComponentSet
from pyomo.core.kernel.variable import variable, variable_dict, variable_list
from pyomo.core.kernel.constraint import constraint, constraint_dict, constraint_list
from pyomo.core.kernel.objective import objective, objective_dict, objective_list
from pyomo.core.kernel.expression import expression, expression_dict, expression_list
from pyomo.core.kernel.block import block, block_dict, block_list
from pyomo.core.kernel.suffix import suffix


class TestComponentSet(unittest.TestCase):
    _components = [
        variable(),
        variable_dict(),
        variable_list(),
        constraint(),
        constraint_dict(),
        constraint_list(),
        objective(),
        objective_dict(),
        objective_list(),
        expression(),
        expression_dict(),
        expression_list(),
        block(),
        block_dict(),
        block_list(),
        suffix(),
    ]

    def test_pickle(self):
        c = ComponentSet()
        self.assertEqual(len(c), 0)
        cup = pickle.loads(pickle.dumps(c))
        self.assertIsNot(cup, c)
        self.assertEqual(len(cup), 0)

        v = variable()
        c.add(v)
        self.assertEqual(len(c), 1)
        self.assertTrue(v in c)
        cup = pickle.loads(pickle.dumps(c))
        vup = cup.pop()
        cup.add(vup)
        self.assertIsNot(cup, c)
        self.assertIsNot(vup, v)
        self.assertEqual(len(cup), 1)
        self.assertTrue(vup in cup)
        self.assertEqual(vup.parent, None)

        b = block()
        V = b.V = variable_list()
        b.V.append(v)
        b.c = c
        self.assertEqual(len(c), 1)
        self.assertTrue(v in c)
        self.assertIs(v.parent, b.V)
        self.assertIs(V.parent, b)
        self.assertIs(b.parent, None)
        bup = pickle.loads(pickle.dumps(b))
        Vup = bup.V
        vup = Vup[0]
        cup = bup.c
        self.assertIsNot(cup, c)
        self.assertIsNot(vup, v)
        self.assertIsNot(Vup, V)
        self.assertIsNot(bup, b)
        self.assertEqual(len(cup), 1)
        self.assertTrue(vup in cup)
        self.assertIs(vup.parent, Vup)
        self.assertIs(Vup.parent, bup)
        self.assertIs(bup.parent, None)

        self.assertEqual(len(c), 1)
        self.assertTrue(v in c)

    def test_init(self):
        cset = ComponentSet()
        cset = ComponentSet(self._components)
        with self.assertRaises(TypeError):
            cset = ComponentSet(*self._components)

    def test_type(self):
        cset = ComponentSet()
        self.assertTrue(isinstance(cset, collections.abc.Set))
        self.assertTrue(isinstance(cset, collections.abc.MutableSet))
        self.assertTrue(issubclass(type(cset), collections.abc.Set))
        self.assertTrue(issubclass(type(cset), collections.abc.MutableSet))

    def test_str(self):
        cset = ComponentSet()
        self.assertEqual(str(cset), "ComponentSet([])")
        cset.update(self._components)
        str(cset)

    def test_len(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        cset.update(self._components)
        self.assertEqual(len(cset), len(self._components))
        cset = ComponentSet(self._components)
        self.assertEqual(len(cset), len(self._components))
        self.assertTrue(len(self._components) > 0)

    def test_iter(self):
        cset = ComponentSet()
        self.assertEqual(list(iter(cset)), [])
        cset.update(self._components)
        ids_seen = set()
        for c in cset:
            ids_seen.add(id(c))
        self.assertEqual(ids_seen, set(id(c) for c in self._components))

    def set_add(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        for i, c in enumerate(self._components):
            self.assertTrue(c not in cset)
            cset.add(c)
            self.assertTrue(c in cset)
            self.assertEqual(len(cset), i + 1)
        self.assertEqual(len(cset), len(self._components))
        for c in self._components:
            self.assertTrue(c in cset)
            cset.add(c)
            self.assertTrue(c in cset)
            self.assertEqual(len(cset), len(self._components))

    def test_pop(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        with self.assertRaises(KeyError):
            cset.pop()
        v = variable()
        cset.add(v)
        self.assertTrue(v in cset)
        self.assertEqual(len(cset), 1)
        v_ = cset.pop()
        self.assertIs(v, v_)
        self.assertTrue(v not in cset)
        self.assertEqual(len(cset), 0)

    def test_update(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        cset.update(self._components)
        self.assertEqual(len(cset), len(self._components))
        for c in self._components:
            self.assertTrue(c in cset)

    def test_clear(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        cset.update(self._components)
        self.assertEqual(len(cset), len(self._components))
        cset.clear()
        self.assertEqual(len(cset), 0)

    def test_remove(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        cset.update(self._components)
        self.assertEqual(len(cset), len(self._components))
        for i, c in enumerate(self._components):
            cset.remove(c)
            self.assertEqual(len(cset), len(self._components) - (i + 1))
        for c in self._components:
            self.assertTrue(c not in cset)
            with self.assertRaises(KeyError):
                cset.remove(c)

    def test_discard(self):
        cset = ComponentSet()
        self.assertEqual(len(cset), 0)
        cset.update(self._components)
        self.assertEqual(len(cset), len(self._components))
        for i, c in enumerate(self._components):
            cset.discard(c)
            self.assertEqual(len(cset), len(self._components) - (i + 1))
        for c in self._components:
            self.assertTrue(c not in cset)
            cset.discard(c)

    def test_isdisjoint(self):
        cset1 = ComponentSet()
        cset2 = ComponentSet()
        self.assertTrue(cset1.isdisjoint(cset2))
        self.assertTrue(cset2.isdisjoint(cset1))
        v = variable()
        cset1.add(v)
        self.assertTrue(cset1.isdisjoint(cset2))
        self.assertTrue(cset2.isdisjoint(cset1))
        cset2.add(v)
        self.assertFalse(cset1.isdisjoint(cset2))
        self.assertFalse(cset2.isdisjoint(cset1))

    def test_misc_set_ops(self):
        v1 = variable()
        cset1 = ComponentSet([v1])
        v2 = variable()
        cset2 = ComponentSet([v2])
        cset3 = ComponentSet([v1, v2])
        empty = ComponentSet([])
        self.assertEqual(cset1 | cset2, cset3)
        self.assertEqual((cset1 | cset2) - cset3, empty)
        self.assertEqual(cset1 ^ cset2, cset3)
        self.assertEqual(cset1 ^ cset3, cset2)
        self.assertEqual(cset2 ^ cset3, cset1)
        self.assertEqual(cset1 & cset2, empty)
        self.assertEqual(cset1 & cset3, cset1)
        self.assertEqual(cset2 & cset3, cset2)

    def test_eq(self):
        cset1 = ComponentSet()
        self.assertEqual(cset1, set())
        self.assertTrue(cset1 == set())
        self.assertNotEqual(cset1, list())
        self.assertFalse(cset1 == list())
        self.assertNotEqual(cset1, tuple())
        self.assertFalse(cset1 == tuple())
        self.assertNotEqual(cset1, dict())
        self.assertFalse(cset1 == dict())

        cset1.update(self._components)
        self.assertNotEqual(cset1, set())
        self.assertFalse(cset1 == set())
        self.assertNotEqual(cset1, list())
        self.assertFalse(cset1 == list())
        self.assertNotEqual(cset1, tuple())
        self.assertFalse(cset1 == tuple())
        self.assertNotEqual(cset1, dict())
        self.assertFalse(cset1 == dict())

        self.assertTrue(cset1 == cset1)
        self.assertEqual(cset1, cset1)

        cset2 = ComponentSet(self._components)
        self.assertTrue(cset2 == cset1)
        self.assertFalse(cset2 != cset1)
        self.assertEqual(cset2, cset1)
        self.assertTrue(cset1 == cset2)
        self.assertFalse(cset1 != cset2)
        self.assertEqual(cset1, cset2)

        cset2.remove(self._components[0])
        self.assertFalse(cset2 == cset1)
        self.assertTrue(cset2 != cset1)
        self.assertNotEqual(cset2, cset1)
        self.assertFalse(cset1 == cset2)
        self.assertTrue(cset1 != cset2)
        self.assertNotEqual(cset1, cset2)

        cset2.add(variable())
        self.assertFalse(cset2 == cset1)
        self.assertTrue(cset2 != cset1)
        self.assertNotEqual(cset2, cset1)
        self.assertFalse(cset1 == cset2)
        self.assertTrue(cset1 != cset2)
        self.assertNotEqual(cset1, cset2)


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