"""
Tests for the omni.graph.tutorials.ExtendedTypes and omni.graph.tutorials.ExtendedTypesPy nodes
"""
import omni.graph.core as og
import omni.graph.core.tests as ogts


class TestTutorialExtendedTypes(ogts.OmniGraphTestCase):
    """Extended attribute type tests require multiple nodes, not supported in the .ogn test framework"""

    # ----------------------------------------------------------------------
    async def _test_tutorial_extended_attributes_node(self, test_node_type_name: str):
        """Test basic operation of the tutorial node containing extended attributes"""

        # Set up a graph that creates full type resolution for simple and array types of the extended attribute node
        #
        #    SimpleIn ----=> Extended1 ----=> SimpleOut
        #              \ /              \ /
        #               X                X
        #              / \              / \
        #    ArrayIn  ----=> Extended2 ----=> ArrayOut
        #
        keys = og.Controller.Keys
        (_, [simple_node, _, extended_node_1, extended_node_2, array_node, _], _, _) = og.Controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("SimpleIn", "omni.graph.tutorials.SimpleData"),
                    ("SimpleOut", "omni.graph.tutorials.SimpleData"),
                    ("Extended1", test_node_type_name),
                    ("Extended2", test_node_type_name),
                    ("ArrayIn", "omni.graph.tutorials.ArrayData"),
                    ("ArrayOut", "omni.graph.tutorials.ArrayData"),
                ],
                keys.CONNECT: [
                    ("SimpleIn.outputs:a_float", "Extended1.inputs:floatOrToken"),
                    ("Extended1.outputs:doubledResult", "SimpleOut.inputs:a_float"),
                    ("ArrayIn.outputs:result", "Extended1.inputs:toNegate"),
                    ("Extended1.outputs:negatedResult", "ArrayOut.inputs:original"),
                    ("SimpleIn.outputs:a_token", "Extended2.inputs:floatOrToken"),
                    ("Extended2.outputs:doubledResult", "SimpleOut.inputs:a_token"),
                    ("ArrayIn.outputs:negativeValues", "Extended2.inputs:toNegate"),
                    ("Extended2.outputs:negatedResult", "ArrayOut.inputs:gates"),
                ],
                keys.SET_VALUES: [
                    ("SimpleIn.inputs:a_float", 5.0),
                    ("SimpleIn.inputs:a_token", "hello"),
                    ("ArrayIn.inputs:multiplier", 2.0),
                    ("ArrayIn.inputs:original", [5.0, -6.0]),
                    ("ArrayIn.inputs:gates", [False, True]),
                ],
            },
        )
        await og.Controller.evaluate()

        # Check that the inputs into the extended type nodes are correct
        self.assertEqual(6.0, og.Controller.get(("outputs:a_float", simple_node)))
        self.assertEqual("world", og.Controller.get(("outputs:a_token", simple_node)))
        self.assertCountEqual([5.0, -12.0], og.Controller.get(("outputs:result", array_node)))
        self.assertCountEqual([False, True], og.Controller.get(("outputs:negativeValues", array_node)))

        # Check the extended simple value outputs
        self.assertEqual(12.0, og.Controller.get(("outputs:doubledResult", extended_node_1)))
        self.assertEqual("worldworld", og.Controller.get(("outputs:doubledResult", extended_node_2)))

        # Check the extended array value outputs
        self.assertCountEqual([-5.0, 12.0], og.Controller.get(("outputs:negatedResult", extended_node_1)))
        self.assertCountEqual([True, False], og.Controller.get(("outputs:negatedResult", extended_node_2)))

    # ----------------------------------------------------------------------
    async def test_tutorial_extended_attributes_node_cpp(self):
        """Test basic operation of the C++ tutorial node containing extended attributes."""
        await self._test_tutorial_extended_attributes_node("omni.graph.tutorials.ExtendedTypes")

    # ----------------------------------------------------------------------
    async def test_tutorial_extended_attributes_node_python(self):
        """Test basic operation of the Python tutorial node containing extended attributes."""
        await self._test_tutorial_extended_attributes_node("omni.graph.tutorials.ExtendedTypesPy")

    # ----------------------------------------------------------------------
    async def _test_tutorial_extended_attributes_tuples(self, test_node_type_name: str):
        """Test basic operation of the tutorial node containing extended attributes on its tuple-accepting attributes"""

        # Set up a graph that creates full type resolution for the tuple types of the extended attribute node with
        # two different resolved types.
        #
        keys = og.Controller.Keys
        (
            _,
            [tuple_node, _, extended_node_1, extended_node_2, tuple_array_node, _, simple_node, _],
            _,
            _,
        ) = og.Controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("TupleIn", "omni.tutorials.TupleData"),
                    ("TupleOut", "omni.tutorials.TupleData"),
                    ("Extended1", test_node_type_name),
                    ("Extended2", test_node_type_name),
                    ("TupleArrayIn", "omni.graph.tutorials.TupleArrays"),
                    ("TupleArrayOut", "omni.graph.tutorials.TupleArrays"),
                    ("SimpleIn", "omni.graph.tutorials.SimpleData"),
                    ("SimpleOut", "omni.graph.tutorials.SimpleData"),
                ],
                keys.CONNECT: [
                    ("TupleIn.outputs:a_int2", "Extended1.inputs:tuple"),
                    ("Extended1.outputs:tuple", "TupleOut.inputs:a_int2"),
                    ("TupleArrayIn.inputs:a", "Extended1.inputs:flexible"),
                    ("Extended1.outputs:flexible", "TupleArrayOut.inputs:a"),
                    ("TupleIn.outputs:a_float3", "Extended2.inputs:tuple"),
                    ("Extended2.outputs:tuple", "TupleOut.inputs:a_float3"),
                    ("SimpleIn.outputs:a_token", "Extended2.inputs:flexible"),
                    ("Extended2.outputs:flexible", "SimpleOut.inputs:a_token"),
                ],
                keys.SET_VALUES: [
                    ("TupleIn.inputs:a_int2", [4, 2]),
                    ("TupleIn.inputs:a_float3", (4.0, 10.0, 2.0)),
                    ("TupleArrayIn.inputs:a", [[2.0, 3.0, 7.0], [21.0, 14.0, 6.0]]),
                    ("TupleArrayIn.inputs:b", [[21.0, 14.0, 6.0], [2.0, 3.0, 7.0]]),
                    ("SimpleIn.inputs:a_token", "hello"),
                ],
            },
        )
        await og.Controller.evaluate()

        # Check that the inputs into the extended type nodes are correct
        self.assertEqual("world", og.Controller.get(("outputs:a_token", simple_node)))
        self.assertCountEqual([5, 3], og.Controller.get(("outputs:a_int2", tuple_node)))
        self.assertCountEqual([5.0, 11.0, 3.0], og.Controller.get(("outputs:a_float3", tuple_node)))
        self.assertCountEqual([126.0, 126.0], og.Controller.get(("outputs:result", tuple_array_node)))

        # Check the resulting values from resolving the "any" type to tuples
        self.assertCountEqual([-5, -3], og.Controller.get(("outputs:tuple", extended_node_1)))
        self.assertCountEqual([-5.0, -11.0, -3.0], og.Controller.get(("outputs:tuple", extended_node_2)))

        # Check the resulting values from resolving the flexible type as both of its types
        self.assertEqual("dlrow", og.Controller.get(("outputs:flexible", extended_node_2)))
        list_expected = [[-2.0, -3.0, -7.0], [-21.0, -14.0, -6.0]]
        list_computed = og.Controller.get(("outputs:flexible", extended_node_1))
        for expected, computed in zip(list_expected, list_computed):
            self.assertCountEqual(expected, computed)

    # ----------------------------------------------------------------------
    async def test_tutorial_extended_attributes_tuples_cpp(self):
        """Test basic operation of the C++ tutorial node containing extended attributes."""
        await self._test_tutorial_extended_attributes_tuples("omni.graph.tutorials.ExtendedTypes")

    # ----------------------------------------------------------------------
    async def test_tutorial_extended_attributes_tuples_python(self):
        """Test basic operation of the Python tutorial node containing extended attributes."""
        await self._test_tutorial_extended_attributes_tuples("omni.graph.tutorials.ExtendedTypesPy")
