#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  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.core.base.indexed_component import normalize_index
from pyomo.core.base.indexed_component_slice import IndexedComponent_slice
from pyomo.core.base.global_set import UnindexedComponent_set
import pyomo.environ as pyo
import pyomo.dae as dae
from pyomo.common.collections import ComponentSet, ComponentMap
from pyomo.util.slices import (
        get_component_call_stack,
        slice_component_along_sets,
        replace_indices,
        get_location_set_map,
        )


class TestGetComponentCallStack(unittest.TestCase):

    get_attribute = IndexedComponent_slice.get_attribute
    get_item = IndexedComponent_slice.get_item

    def assertSameStack(self, stack1, stack2):
        for (call1, arg1), (call2, arg2) in zip(stack1, stack2):
            self.assertIs(call1, call2)
            self.assertEqual(arg1, arg2)

    def assertCorrectStack(self, comp, pred_stack, context=None):
        act_stack = get_component_call_stack(comp, context=None)
        self.assertSameStack(pred_stack, act_stack)

    def model(self):

        m = pyo.ConcreteModel()
        m.s1 = pyo.Set(initialize=[1,2,3])
        m.s2 = pyo.Set(initialize=[('a',1),('b',2)])

        m.v0 = pyo.Var()
        m.v1 = pyo.Var(m.s1)
        m.v2 = pyo.Var(m.s2)

        @m.Block()
        def b(b):

            @b.Block(m.s1)
            def b1(b):
                b.v0 = pyo.Var()
                b.v1 = pyo.Var(m.s1)

                @b.Block()
                def b(b):
                    b.v0 = pyo.Var()
                    b.v2 = pyo.Var(m.s2)

                    @b.Block(m.s1, m.s2)
                    def b2(b,i1,i2,i3=None):
                        # Optional i3 is to allow for the case
                        # when normalize_index.flatten is False
                        b.v0 = pyo.Var()
                        b.v1 = pyo.Var(m.s1)

                        @b.Block()
                        def b(b):
                            b.v0 = pyo.Var()
                            b.v2 = pyo.Var(m.s2)
        return m

    def test_no_context(self):
        m = self.model()
        
        comp = m.v1[1]
        pred_stack = [
                (self.get_item, 1),
                (self.get_attribute, 'v1'), 
                ]
        self.assertCorrectStack(comp, pred_stack)

        comp = m.v1
        pred_stack = [
                (self.get_attribute, 'v1'),
                ]
        self.assertCorrectStack(comp, pred_stack)

        comp = m.b.b1[1].b.b2
        pred_stack = [
                (self.get_attribute, 'b2'),
                (self.get_attribute, 'b'),
                (self.get_item, 1),
                (self.get_attribute, 'b1'),
                (self.get_attribute, 'b'),
                ]
        self.assertCorrectStack(comp, pred_stack)

        comp = m.b.b1[1].b.b2[1,'a',1]
        pred_stack = [
                (self.get_item, (1,'a',1)),
                (self.get_attribute, 'b2'),
                (self.get_attribute, 'b'),
                (self.get_item, 1),
                (self.get_attribute, 'b1'),
                (self.get_attribute, 'b'),
                ]
        self.assertCorrectStack(comp, pred_stack)

        normalize_index.flatten = False
        comp = m.b.b1[1].b.b2[1,('a',1)]
        # NOTE: This fails without the parenthesis around ('a',1)
        # as comp.index() is None. The actual stack just starts one
        # layer higher...
        # It is confusing that comp can even be generated in this case;
        # shouldn't (1,'a',1) be an invalid index?
        pred_stack = [
                (self.get_item, (1,('a',1))),
                (self.get_attribute, 'b2'),
                (self.get_attribute, 'b'),
                (self.get_item, 1),
                (self.get_attribute, 'b1'),
                (self.get_attribute, 'b'),
                ]
        self.assertCorrectStack(comp, pred_stack)
        normalize_index.flatten = True

    def test_from_block(self):
        m = self.model()

        comp = m.v0
        pred_stack = [
                (self.get_attribute, 'v0'),
                ]
        self.assertCorrectStack(comp, pred_stack, context=m)

        comp = m.b.b1[2].b.b2[1,'a',1]
        pred_stack = [
                (self.get_item, (1,'a',1)),
                (self.get_attribute, 'b2'),
                (self.get_attribute, 'b'),
                (self.get_item, 2),
                (self.get_attribute, 'b1'),
                (self.get_attribute, 'b'),
                ]
        self.assertCorrectStack(comp, pred_stack, context=m)

        comp = m.b.b1[2].b.b2[1,'a',1].b.v2['b',2]
        pred_stack = [
                (self.get_item, ('b',2)),
                (self.get_attribute, 'v2'),
                (self.get_attribute, 'b'),
                (self.get_item, (1,'a',1)),
                (self.get_attribute, 'b2'),
                (self.get_attribute, 'b'),
                (self.get_item, 2),
                ]
        self.assertCorrectStack(comp, pred_stack, context=m.b.b1)

        comp = m.b.b1[2].b.b2
        pred_stack = [
                (self.get_attribute, 'b2'),
                (self.get_attribute, 'b'),
                (self.get_item, 2),
                ]
        self.assertCorrectStack(comp, pred_stack, context=m.b.b1)

        comp = m.b.b1[2]
        pred_stack = [
                (self.get_item, 2),
                ]
        self.assertCorrectStack(comp, pred_stack, context=m.b.b1)

        comp = m.b.b1
        act_stack = get_component_call_stack(comp, context=m.b.b1)
        self.assertEqual(len(act_stack), 0)

    def test_from_blockdata(self):

        m = self.model()

        context = m.b.b1[3].b.b2[2,'b',2]
        comp = m.b.b1[3].b.b2[2,'b',2].b
        pred_stack = [
                (self.get_attribute, 'b'),
                ]
        self.assertCorrectStack(comp, pred_stack, context=context)

        comp = m.b.b1[3].b.b2[2,'b',2].b.v2['a',1]
        pred_stack = [
                (self.get_item, ('a',1)),
                (self.get_attribute, 'v2'),
                (self.get_attribute, 'b'),
                ]
        self.assertCorrectStack(comp, pred_stack, context=context)

        context = m.b.b1[3]
        comp = m.b.b1[3]
        act_stack = get_component_call_stack(comp, context=context)
        self.assertEqual(len(act_stack), 0)

class TestGetLocationAndReplacement(unittest.TestCase):

    def model(self):

        m = pyo.ConcreteModel()

        m.time = pyo.Set(initialize=[1,2,3])
        m.space = dae.ContinuousSet(initialize=[0,2])
        m.comp = pyo.Set(initialize=['a','b'])
        m.d_2 = pyo.Set(initialize=[('a',1),('b',2)])
        m.d_none = pyo.Set(initialize=[('c',1,10), ('d',3)], dimen=None)

        m.b0 = pyo.Block()
        m.b1 = pyo.Block(m.time)
        m.b2 = pyo.Block(m.time, m.space)
        m.b0d2 = pyo.Block(m.d_2)
        m.b1d2 = pyo.Block(m.time, m.d_2)
        m.b2d2 = pyo.Block(m.time, m.d_2, m.space)
        m.b3d2 = pyo.Block(m.time, m.d_2, m.space, m.time)

        m.b0dn = pyo.Block(m.d_none)
        m.b1dn = pyo.Block(m.time, m.d_none)
        m.b1dnd2 = pyo.Block(m.time, m.d_none, m.d_2)
        m.dnd2b1 = pyo.Block(m.d_none, m.d_2, m.time)
        m.b2dn = pyo.Block(m.time, m.space, m.d_none)
        m.b2dnd2 = pyo.Block(m.time, m.d_none, m.d_2, m.space)
        m.b3dn = pyo.Block(m.time, m.d_2, m.d_none, m.space, m.d_2)

        m.dn2 = pyo.Block(m.time, m.d_none, m.d_2, m.d_none, m.d_2)

        return m

    def assertSameMap(self, pred, act):
        self.assertEqual(len(pred), len(act))
        for k, v in pred.items():
            self.assertIn(k, act)
            self.assertIs(pred[k], act[k])
        
    def test_one_to_one(self):
        m = self.model()

        index_set = m.b0.index_set()
        index = None
        pred_map = {0: UnindexedComponent_set}
        location_set_map = get_location_set_map(index, index_set)
        self.assertEqual(pred_map, location_set_map)

        index_set = m.b1.index_set()
        index = 1
        pred_map = {
                0: m.time,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b2.index_set()
        index = (1,2)
        pred_map = {
                0: m.time,
                1: m.space,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

    def test_1d_replacement(self):
        m = self.model()

        # Unindexed
        index_set = m.b0.index_set()
        index = None
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet()
        pred_index = (None,)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(pred_index, new_index)

        # One one-d set
        index_set = m.b1.index_set()
        index = 1
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet()
        pred_index = (1,)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(pred_index, new_index)

        sets = ComponentSet((m.time,))
        pred_index = (slice(None),)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(pred_index, new_index)

        # Two one-d sets
        index_set = m.b2.index_set()
        index = (1,2)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet()
        pred_index = (1,2)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(pred_index, new_index)

        sets = ComponentSet((m.space,))
        pred_index = (1, slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(pred_index, new_index)

        sets = ComponentSet((m.space, m.time))
        pred_index = (slice(None), slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(pred_index, new_index)

    def test_multi_dim(self):
        m = self.model()
        index_set = m.b0d2.index_set()
        index = ('a',1)
        pred_map = {
                0: m.d_2,
                1: m.d_2,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        normalize_index.flatten = False
        # A single multi-dimensional index will behave the same when
        # normalize_index.flatten is False. This is probably unexpected
        # Should we just fail in this case?
        # Note that the index (('a',1),) will not work.
        index_set = m.b0d2.index_set()
        index = ('a',1)
        pred_map = {
                0: m.d_2,
                1: m.d_2,
                }
        with self.assertRaises(ValueError) as cm:
            location_set_map = get_location_set_map(index, index_set)
        self.assertIn('normalize_index.flatten', str(cm.exception))
        normalize_index.flatten = True

        index_set = m.b1d2.index_set()
        index = (1, 'a', 1)
        pred_map = {
                0: m.time,
                1: m.d_2,
                2: m.d_2,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b2d2.index_set()
        index = (1, 'a', 1, 2)
        pred_map = {
                0: m.time,
                1: m.d_2,
                2: m.d_2,
                3: m.space,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b3d2.index_set()
        index = (1, 'a', 1, 2, 1)
        pred_map = {
                0: m.time,
                1: m.d_2,
                2: m.d_2,
                3: m.space,
                4: m.time,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

    def test_multi_dim_replacement(self):
        m = self.model()

        # One two-dimen set
        index_set = m.b0d2.index_set()
        index = ('a',1)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet()
        pred_index = ('a', 1)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.d_2,))
        pred_index = (slice(None), slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One two-dimen, one one-dimen
        index_set = m.b1d2.index_set()
        index = (1, 'a', 1)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet()
        pred_index = (1, 'a', 1)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.d_2,))
        pred_index = (1, slice(None), slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One two-dimen, two one-dimen
        index_set = m.b2d2.index_set()
        index = (1, 'a', 1, 2)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_2, m.time))
        pred_index = (slice(None), slice(None), slice(None), 2)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.space,))
        pred_index = (1, 'a', 1, slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One two-dimen, three one-dimen
        index_set = m.b3d2.index_set()
        index = (1, 'a', 1, 2, 1)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.time,))
        pred_index = (slice(None), 'a', 1, 2, slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

    def test_dimen_none(self):
        m = self.model()

        index_set = m.b0dn.index_set()
        index = ('c',1,10)
        pred_map = {
                0: m.d_none,
                1: m.d_none,
                2: m.d_none,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b1dn.index_set()
        index = (1, 'c', 1, 10)
        pred_map = {
                0: m.time,
                1: m.d_none,
                2: m.d_none,
                3: m.d_none,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b1dnd2.index_set()
        index = (1, 'c', 1, 10, 'a', 1)
        pred_map = {
                0: m.time,
                1: m.d_none,
                2: m.d_none,
                3: m.d_none,
                4: m.d_2,
                5: m.d_2,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b2dn.index_set()
        index = (1, 0, 'd', 3)
        pred_map = {
                0: m.time,
                1: m.space,
                2: m.d_none,
                3: m.d_none,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b2dnd2.index_set()
        index = (1, 'c', 1, 10, 'b', 2, 0)
        pred_map = {
                0: m.time,
                1: m.d_none,
                2: m.d_none,
                3: m.d_none,
                4: m.d_2,
                5: m.d_2,
                6: m.space,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.dnd2b1.index_set()
        index = ('d', 3, 'b', 2, 1)
        pred_map = {
                0: m.d_none,
                1: m.d_none,
                2: m.d_2,
                3: m.d_2,
                4: m.time,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.b3dn.index_set()
        index = (1, 'a', 1, 'd', 3, 0, 'b', 2)
        pred_map = {
                0: m.time,
                1: m.d_2,
                2: m.d_2,
                3: m.d_none,
                4: m.d_none,
                5: m.space,
                6: m.d_2,
                7: m.d_2,
                }
        location_set_map = get_location_set_map(index, index_set)
        self.assertSameMap(pred_map, location_set_map)

        index_set = m.dn2.index_set()
        index = (1, 'c', 1, 10, 'b', 2, 'd', 3, 'a', 1)
        with self.assertRaises(RuntimeError) as cm:
            location_set_map = get_location_set_map(index, index_set)
        self.assertIn('multiple sets of dimen==None', str(cm.exception))

    def test_dimen_none_replacement(self):
        m = self.model()

        # Single set of dimen None
        index_set = m.b0dn.index_set()
        index = ('c',1,10)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_none,))
        pred_index = (Ellipsis,)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet()
        pred_index = ('c', 1, 10)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One None-dimen, one one-dimen
        index_set = m.b1dn.index_set()
        index = (1, 'c', 1, 10)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_none,))
        pred_index = (1, Ellipsis)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.d_none, m.time))
        pred_index = (slice(None), Ellipsis)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One None-dimen, one one-dimen, one two-dimen
        index_set = m.b1dnd2.index_set()
        index = (1, 'c', 1, 10, 'a', 1)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_none, m.time, m.d_2))
        pred_index = (slice(None), Ellipsis, slice(None), slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.time, m.d_2))
        pred_index = (slice(None), 'c', 1, 10, slice(None), slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.d_2,))
        pred_index = (1, 'c', 1, 10, slice(None), slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One None-dimen, one two-dimen
        index_set = m.b2dn.index_set()
        index = (1, 0, 'd', 3)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet()
        pred_index = (1, 0, 'd', 3)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        sets = ComponentSet((m.d_none,))
        pred_index = (1, 0, Ellipsis)
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # One None-dimen, one two-dimen, two one-dimen
        index_set = m.b2dnd2.index_set()
        index = (1, 'c', 1, 10, 'b', 2, 0)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_none, m.space))
        pred_index = (1, Ellipsis, 'b', 2, slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # None-dimen set comes first
        index_set = m.dnd2b1.index_set()
        index = ('d', 3, 'b', 2, 1)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_none, m.time))
        pred_index = (Ellipsis, 'b', 2, slice(None))
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

        # None-dimen, plus 2-d set shows up twice
        index_set = m.b3dn.index_set()
        index = (1, 'a', 1, 'd', 3, 0, 'b', 2)
        location_set_map = get_location_set_map(index, index_set)
        sets = ComponentSet((m.d_none, m.d_2))
        pred_index = (
                1,
                slice(None),
                slice(None),
                Ellipsis,
                0,
                slice(None),
                slice(None),
                )
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)

    def test_replacement_normalize_flatten_false(self):
        m = self.model()

        # index of m.b2d2
        index = (1, 0, ('a', 1))
        location_set_map = {
                0: m.time,
                1: m.space,
                2: m.d_2,
                }
        sets = ComponentSet((m.d_2,))
        pred_index = (1, 0, slice(None))
        # If you can assemble the proper location_set_map 
        # (get_location_set_map will not work), this function
        # should work when normalize_index.flatten is False.
        normalize_index.flatten = False
        new_index = replace_indices(index, location_set_map, sets)
        self.assertEqual(new_index, pred_index)
        normalize_index.flatten = True


class SliceComponentDataAlongSets(unittest.TestCase):
    def model(self):
        m = pyo.ConcreteModel()

        m.time = pyo.Set(initialize=[1,2,3])
        m.space = dae.ContinuousSet(initialize=[0,2])
        m.comp = pyo.Set(initialize=['a','b'])
        m.d_2 = pyo.Set(initialize=[('a',1),('b',2)])
        m.d_none = pyo.Set(initialize=[('c',1,10), ('d',3)], dimen=None)

        @m.Block()
        def b(b):
            b.v0 = pyo.Var()
            b.v1 = pyo.Var(m.time)
            b.v2 = pyo.Var(m.time, m.space)

            @b.Block(m.time, m.space)
            def b2(b2):
                b2.v0 = pyo.Var()
                b2.v1 = pyo.Var(m.comp)
                b2.v2 = pyo.Var(m.time, m.comp)
                b2.vn = pyo.Var(m.time, m.d_none, m.d_2)
    
            @b.Block(m.d_none, m.d_2)
            def bn(bn):
                bn.v0 = pyo.Var()
                bn.v2 = pyo.Var(m.time, m.space)
                bn.v3 = pyo.Var(m.time, m.space, m.time)
                bn.vn = pyo.Var(m.time, m.d_none, m.d_2)
        
        return m

    def test_with_tuple_of_sets(self):
        m = pyo.ConcreteModel()
        m.s1 = pyo.Set(initialize=[1, 2, 3])
        m.s2 = pyo.Set(initialize=[1, 2, 3])
        m.v = pyo.Var(m.s1, m.s2)
        sets = (m.s1,)
        slice_ = slice_component_along_sets(m.v[1, 2], sets)

        # These tests are essentially the same, but we run the CUID test
        # first as it gives a nicer error message.
        self.assertEqual(str(pyo.ComponentUID(slice_)), "v[*,2]")
        self.assertEqual(slice_, m.v[:, 2])

    def test_no_context(self):
        m = self.model()
        
        comp = m.b.v0
        sets = ComponentSet((m.time, m.space))
        _slice = slice_component_along_sets(comp, sets)
        # Use `assertIs` when the "slice" is actually just a component.
        self.assertIs(_slice, m.b.v0)

        comp = m.b.v1[1]
        sets = ComponentSet((m.time, m.space))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.v1[:])

        comp = m.b.v2[1,0]
        sets = ComponentSet((m.time, m.space))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.v2[:,:])

        comp = m.b.b2[1,0].v1['a']
        sets = ComponentSet((m.time, m.space))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.b2[:,:].v1['a'])

        comp = m.b.b2[1,0].v1
        sets = ComponentSet((m.time, m.space))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.b2[:,:].v1)

        comp = m.b.b2[1,0].v2[1,'a']
        sets = ComponentSet((m.time,))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.b2[:,0].v2[:,'a'])

        comp = m.b.bn['c',1,10,'a',1].v3[1,0,1]
        sets = ComponentSet((m.time,))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.bn['c',1,10,'a',1].v3[:,0,:])

    def test_context(self):
        m = self.model()

        comp = m.b.v2[1,0]
        context = m.b
        sets = ComponentSet((m.time,))
        _slice = slice_component_along_sets(comp, sets)
        # Context makes no difference in resulting slice here.
        self.assertEqual(_slice, m.b.v2[:,0])

        comp = m.b.b2[1,0].v2[1,'a']
        context = m.b.b2[1,0]
        sets = ComponentSet((m.time,))
        _slice = slice_component_along_sets(comp, sets, context=context)
        # Here, context does make a difference
        self.assertEqual(_slice, m.b.b2[1,0].v2[:,'a'])

        comp = m.b.b2[1,0].v1['a']
        context = m.b.b2[1,0]
        sets = ComponentSet((m.time,))
        _slice = slice_component_along_sets(comp, sets, context=context)
        self.assertIs(_slice, m.b.b2[1,0].v1['a'])

        sets = ComponentSet((m.comp,))
        _slice = slice_component_along_sets(comp, sets, context=context)
        self.assertEqual(_slice, m.b.b2[1,0].v1[:])

        context = m.b.b2
        sets = ComponentSet((m.time, m.comp))
        _slice = slice_component_along_sets(comp, sets, context=context)
        self.assertEqual(_slice, m.b.b2[:,0].v1[:])

    def test_none_dimen(self):
        m = self.model()

        comp = m.b.b2[1,0].vn
        sets = ComponentSet((m.d_none,))
        _slice = slice_component_along_sets(comp, sets)
        self.assertIs(_slice, m.b.b2[1,0].vn)

        comp = m.b.b2[1,0].vn[1,'d',3,'a',1]
        sets = ComponentSet((m.d_2, m.time))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.b2[:,0].vn[:,'d',3,:,:])

        sets = ComponentSet((m.d_none, m.d_2))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.b2[1,0].vn[1,...,:,:])

        comp = m.b.bn['c',1,10,'a',1].v3[1,0,1]
        sets = ComponentSet((m.d_none, m.time))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.bn[...,'a',1].v3[:,0,:])

        comp = m.b.bn['c',1,10,'a',1].vn[1,'d',3,'b',2]
        sets = ComponentSet((m.d_none,))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.bn[...,'a',1].vn[1,...,'b',2])

        sets = ComponentSet((m.d_none, m.d_2))
        _slice = slice_component_along_sets(comp, sets)
        self.assertEqual(_slice, m.b.bn[...,:,:].vn[1,...,:,:])
        
        comp = m.b.bn['c',1,10,'a',1].vn[1,'d',3,'b',2]
        context = m.b.bn['c',1,10,'a',1]
        sets = ComponentSet((m.d_none,))
        _slice = slice_component_along_sets(comp, sets, context=context)
        self.assertEqual(_slice, m.b.bn['c',1,10,'a',1].vn[1,...,'b',2])

        sets = ComponentSet((m.d_none, m.d_2))
        _slice = slice_component_along_sets(comp, sets, context=context)
        self.assertEqual(_slice, m.b.bn['c',1,10,'a',1].vn[1,...,:,:])

        context = m.b.bn
        sets = ComponentSet((m.d_none,))
        _slice = slice_component_along_sets(comp, sets, context=context)
        self.assertEqual(_slice, m.b.bn[...,'a',1].vn[1,...,'b',2])


class TestCUID(unittest.TestCase):

    def test_cuid_of_slice(self):
        m = pyo.ConcreteModel()
        m.s1 = pyo.Set(initialize=["a", "b"])
        m.s2 = pyo.Set(initialize=["c", "d"])
        m.b = pyo.Block(m.s1)
        for i in m.s1:
            m.b[i].v = pyo.Var(m.s2)
        slice_ = slice_component_along_sets(m.b["a"].v["c"], ComponentSet((m.s1,)))
        cuid = pyo.ComponentUID(slice_)
        self.assertEqual(str(cuid), "b[*].v[c]")

        slice_ = slice_component_along_sets(
            m.b["a"].v[("c",)], ComponentSet((m.s1,))
        )
        cuid = pyo.ComponentUID(slice_)
        self.assertEqual(str(cuid), "b[*].v[c]")


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