"""Action Graph Evaluation Tests, Part 2"""

from enum import Enum, auto

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from omni.graph.action import get_interface
from pxr import Gf


# ======================================================================
class TestActionGraphEvaluation(ogts.OmniGraphTestCase):
    """Tests action graph evaluator functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    async def setUp(self):
        """Set up  test environment, to be torn down when done"""
        await super().setUp()
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

        # ----------------------------------------------------------------------

    async def test_latent_fan_out(self):
        """Test latent nodes when part of parallel evaluation"""
        #                                              +------------+
        #                                        +---->|TickCounterA|
        #                                        |     +------------+
        #                                        |
        #                               +--------++    +----------+
        #                             +-> TickA   +--->|FinishedA |
        #                             | +---------+    +----------+
        # +---------+   +-----------+ |
        # |OnImpulse+-->|TickCounter+-+
        # +---------+   +-----------+ |
        #                             | +---------+    +----------+
        #                             +>| TickB   +--->|FinishedB |
        #                               +--------++    +----------+
        #                                        |
        #                                        |     +------------+
        #                                        +---->|TickCounterB|
        #                                              +------------+
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("TickA", "omni.graph.action.Countdown"),
                    ("TickB", "omni.graph.action.Countdown"),
                    ("TickCounter", "omni.graph.action.Counter"),
                    ("TickCounterA", "omni.graph.action.Counter"),
                    ("TickCounterB", "omni.graph.action.Counter"),
                    ("FinishCounterA", "omni.graph.action.Counter"),
                    ("FinishCounterB", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "TickCounter.inputs:execIn"),
                    ("TickCounter.outputs:execOut", "TickA.inputs:execIn"),
                    ("TickCounter.outputs:execOut", "TickB.inputs:execIn"),
                    ("TickA.outputs:tick", "TickCounterA.inputs:execIn"),
                    ("TickB.outputs:tick", "TickCounterB.inputs:execIn"),
                    ("TickA.outputs:finished", "FinishCounterA.inputs:execIn"),
                    ("TickB.outputs:finished", "FinishCounterB.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("TickA.inputs:duration", 2),
                    ("TickB.inputs:duration", 2),
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("OnImpulse.state:enableImpulse", True),
                ],
            },
        )

        (_, _, _, tick_counter, tick_counter_a, tick_counter_b, finish_counter_a, finish_counter_b) = nodes

        def check_counts(t_def, t_a, t_b, f_a, f_b):
            for node, expected in (
                (tick_counter, t_def),
                (tick_counter_a, t_a),
                (tick_counter_b, t_b),
                (finish_counter_a, f_a),
                (finish_counter_b, f_b),
            ):
                count = og.Controller.get(controller.attribute("outputs:count", node))
                self.assertEqual(count, expected, node.get_prim_path())

        await controller.evaluate(graph)
        check_counts(1, 0, 0, 0, 0)
        await controller.evaluate(graph)
        check_counts(1, 1, 1, 0, 0)
        await controller.evaluate(graph)
        check_counts(1, 2, 2, 0, 0)
        await controller.evaluate(graph)
        check_counts(1, 2, 2, 1, 1)

    # ----------------------------------------------------------------------
    async def test_loop_cancel(self):
        """Test loop canceling"""
        # Check that a loop can be canceled.
        # We set up the loop to run for 7 iterations, but cancel when we hit iteration 2.
        #
        #                                      +--------+
        #                       +------------->|COUNTER2|
        #                       |   finish     +--------+
        # +---------+   +----------+    +----------+    +-------+    +---------+
        # | IMPULSE +-->| FOR-LOOP +--->| COMPARE  +--->|BRANCH +--->| COUNTER1|
        # +---------+   +----------+    +----------+    +-------+    +---------+
        #                       ^                              |
        #                       |      cancel                  |
        #                       +------------------------------+
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, for_loop, _, _, count_1, count_2), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("ForLoop", "omni.graph.action.ForLoop"),
                    ("Compare", "omni.graph.nodes.Compare"),
                    ("Branch", "omni.graph.action.Branch"),
                    ("Counter1", "omni.graph.action.Counter"),
                    ("Counter2", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "ForLoop.inputs:execIn"),
                    ("ForLoop.outputs:loopBody", "Branch.inputs:execIn"),
                    ("ForLoop.outputs:index", "Compare.inputs:a"),
                    ("ForLoop.outputs:finished", "Counter2.inputs:execIn"),
                    ("Compare.outputs:result", "Branch.inputs:condition"),
                    ("Branch.outputs:execFalse", "Counter1.inputs:execIn"),
                    ("Branch.outputs:execTrue", "ForLoop.inputs:breakLoop"),
                ],
                keys.SET_VALUES: [
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("Compare.inputs:b", 2, "int"),
                    ("Compare.inputs:operation", "=="),
                    ("ForLoop.inputs:stop", 6),
                ],
            },
        )
        await controller.evaluate(graph)

        # Trigger graph once.
        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: ("OnImpulse.state:enableImpulse", True)})
        await controller.evaluate(graph)

        # Verify the loop body only counted 2 times, and finish once.
        c_1 = og.Controller.get(og.Controller.attribute("outputs:count", count_1))
        c_2 = og.Controller.get(og.Controller.attribute("outputs:count", count_2))
        index = og.Controller.get(og.Controller.attribute("outputs:index", for_loop))
        self.assertEqual(index, 2)
        self.assertEqual(c_1, 2)
        self.assertEqual(c_2, 1)

    # ----------------------------------------------------------------------
    async def test_loop_sideffects(self):
        """Test that nodes in loops can read results of external side effects"""
        # +----------+      +---------+      +-------+       +--------------------+
        # | OnImpulse+----->| ForLoop +----->|  Add  +------>| WritePrimAttribute |
        # +----------+      +---------+      +-------+       +--------------------+
        #                                     ^
        #               +-----------------+   |
        #               |ReadPrimAttribute +---+
        #               +-----------------+

        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                    ("ForLoop1", "omni.graph.action.ForLoop"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: ("/World/Accumulator", {"acc": ("int", 0)}),
                keys.SET_VALUES: [
                    ("Get.inputs:name", "acc"),
                    ("Get.inputs:usePath", True),
                    ("Get.inputs:primPath", "/World/Accumulator"),
                    ("ForLoop1.inputs:start", 0),
                    ("ForLoop1.inputs:stop", 5),
                    ("Set.inputs:name", "acc"),
                    ("Set.inputs:usePath", True),
                    ("Set.inputs:primPath", "/World/Accumulator"),
                    ("OnImpulse.inputs:onlyPlayback", False),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "ForLoop1.inputs:execIn"),
                    ("ForLoop1.outputs:loopBody", "Set.inputs:execIn"),
                    ("Get.outputs:value", "Add.inputs:a"),
                    ("ForLoop1.outputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "Set.inputs:value"),
                ],
            },
        )
        await controller.evaluate(graph)

        # Trigger graph evaluation once.
        # Should result in sum = 0 + 0 + 1 + 2 + 3 + 4.
        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: [("OnImpulse.state:enableImpulse", True)]})
        await controller.evaluate(graph)
        self.assertEqual(prims[0].GetAttribute("acc").Get(), 10)

    # ----------------------------------------------------------------------
    async def test_nested_forloop(self):
        """Test nested ForLoop nodes"""
        keys = og.Controller.Keys
        controller = og.Controller()
        (graph, (_, _, _, _, _, _, write_node), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Const", "omni.graph.nodes.ConstantInt2"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("Branch", "omni.graph.action.Branch"),
                    ("For", "omni.graph.action.ForLoop"),
                    ("For2", "omni.graph.action.ForLoop"),
                    ("Write1", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: (
                    "/World/TestPrim",
                    {
                        "val1": ("int[2]", Gf.Vec2i(1, 1)),
                    },
                ),
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Const.inputs:value", [1, 2]),
                    ("Write1.inputs:name", "val1"),
                    ("Write1.inputs:primPath", "/World/TestPrim"),
                    ("Write1.inputs:usePath", True),
                    ("For.inputs:stop", 3),
                    ("For2.inputs:start", 4),
                    ("For2.inputs:step", 2),
                    ("For2.inputs:stop", 10),
                    ("Branch.inputs:condition", True),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "For.inputs:execIn"),
                    ("For.outputs:loopBody", "Branch.inputs:execIn"),
                    ("Branch.outputs:execTrue", "For2.inputs:execIn"),
                    ("For2.outputs:loopBody", "Write1.inputs:execIn"),
                    ("For2.outputs:value", "Add.inputs:a"),
                    ("Const.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "Write1.inputs:value"),
                ],
            },
        )
        await controller.evaluate(graph)
        context = omni.usd.get_context()
        stage = context.get_stage()

        # For2 should loop over the range [4, 6, 8, 10).
        self.assertEqual(9, write_node.get_compute_count())
        # [1, 2] + 8 = [9, 10].
        self.assertListEqual([9, 10], list(stage.GetAttributeAtPath("/World/TestPrim.val1").Get()))

    # ----------------------------------------------------------------------
    async def test_om_63924(self):
        """Test OM-63924 bug is fixed"""
        # The problem here was that if there was fan in to a node which was
        # computed once and then totally unwound before the other history was
        # processed, there would never be a deferred activation and so the 2nd
        # compute would never happen. Instead we want to only unwind one history
        # at a time to ensure each one is fully evaluated.
        i = 2

        class OnForEachEventPy:
            """Helper Python node that implements ForEach logic"""

            @staticmethod
            def compute(context: og.GraphContext, node: og.Node):
                """Compute method"""
                nonlocal i
                inputs_go = node.get_attribute("inputs:go")
                go_val = og.Controller.get(inputs_go)
                if not go_val:
                    return True
                if i > 0:
                    og.Controller.set(
                        node.get_attribute("outputs:execOut"), og.ExecutionAttributeState.ENABLED_AND_PUSH
                    )
                    og.Controller.set(node.get_attribute("outputs:syncValue"), i)
                    i -= 1
                return True

            @staticmethod
            def get_node_type() -> str:
                """Get node type"""
                return "omni.graph.test.OnForEachEventPy"

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                """Initialize node attributes"""
                node_type.add_input(
                    "inputs:go",
                    "bool",
                    False,
                )
                node_type.add_output("outputs:execOut", "execution", True)
                node_type.add_output("outputs:syncValue", "uint64", True)
                return True

        og.register_node_type(OnForEachEventPy, 1)

        class NoOpPy:
            """Helper Python node that performs no internal operation"""

            @staticmethod
            def compute(context: og.GraphContext, node: og.Node):
                """Compute method"""
                og.Controller.set(node.get_attribute("outputs:execOut"), og.ExecutionAttributeState.ENABLED)
                return True

            @staticmethod
            def get_node_type() -> str:
                """Get node type"""
                return "omni.graph.test.NoOpPy"

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                """Initialize node attributes"""
                node_type.add_input(
                    "inputs:execIn",
                    "execution",
                    True,
                )
                node_type.add_output("outputs:execOut", "execution", True)
                return True

        og.register_node_type(NoOpPy, 1)

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (for_each, _, _, _, _, no_op_2), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("PostProcessDispatcher", "omni.graph.test.OnForEachEventPy"),
                    ("TSA1", "omni.graph.action.SyncGate"),
                    ("TSA0", "omni.graph.action.SyncGate"),
                    ("TestSyncAccum", "omni.graph.action.SyncGate"),
                    ("TestPrimBbox", "omni.graph.test.NoOpPy"),
                    ("NoOpPy2", "omni.graph.test.NoOpPy"),
                ],
                keys.CONNECT: [
                    ("PostProcessDispatcher.outputs:execOut", "TSA0.inputs:execIn"),
                    ("PostProcessDispatcher.outputs:execOut", "TSA1.inputs:execIn"),
                    ("TSA1.outputs:execOut", "TestSyncAccum.inputs:execIn"),
                    ("TSA0.outputs:execOut", "TestPrimBbox.inputs:execIn"),
                    ("TestPrimBbox.outputs:execOut", "TestSyncAccum.inputs:execIn"),
                    ("TestSyncAccum.outputs:execOut", "NoOpPy2.inputs:execIn"),
                    ("PostProcessDispatcher.outputs:syncValue", "TSA1.inputs:syncValue"),
                    ("PostProcessDispatcher.outputs:syncValue", "TSA0.inputs:syncValue"),
                    ("PostProcessDispatcher.outputs:syncValue", "TestSyncAccum.inputs:syncValue"),
                ],
            },
        )
        og.Controller.set(controller.attribute("inputs:go", for_each), True)
        await controller.evaluate(graph)
        # Verify the final sync gate triggered due to being computed 2x.
        exec_out = og.Controller.get(controller.attribute("outputs:execOut", no_op_2))
        self.assertEqual(exec_out, og.ExecutionAttributeState.ENABLED)

    # ----------------------------------------------------------------------
    async def test_retrigger_latent(self):
        """Test that latent nodes can be re-triggered"""
        want_debug = False
        e_state = og.ExecutionAttributeState
        tick_count = 0
        boop_count = 0
        exec_in_latent_count = 0
        max_ticks = 20

        class CancelTickerPy:
            """Helper node type which does latent ticking and can be canceled, and has an independent counter "boop" """

            @staticmethod
            def compute(context: og.GraphContext, node: og.Node):
                """Compute method"""
                nonlocal tick_count
                nonlocal boop_count
                nonlocal exec_in_latent_count
                exec_in = node.get_attribute("inputs:execIn")
                exec_in_val = og.Controller.get(exec_in)
                cancel = node.get_attribute("inputs:cancel")
                cancel_val = og.Controller.get(cancel)
                boop = node.get_attribute("inputs:boop")
                boop_val = og.Controller.get(boop)

                if want_debug:
                    print(f"### {tick_count} execIn={exec_in_val} cancel={cancel_val} boop={boop_val}")

                if cancel_val == e_state.ENABLED:
                    # Finish latent by cancel.
                    og.Controller.set(node.get_attribute("outputs:canceled"), e_state.LATENT_FINISH)
                    self.assertEqual(exec_in_val, e_state.DISABLED)
                    self.assertEqual(boop_val, e_state.DISABLED)
                    tick_count = 0
                    return True
                if exec_in_val == e_state.ENABLED:
                    self.assertEqual(cancel_val, e_state.DISABLED)
                    self.assertEqual(boop_val, e_state.DISABLED)
                    if tick_count > 0:
                        # execIn triggered while in latent - should not be possible.
                        exec_in_latent_count += 1
                    else:
                        og.Controller.set(node.get_attribute("outputs:tick"), e_state.LATENT_PUSH)
                        return True

                # We are ticking.
                self.assertEqual(cancel_val, e_state.DISABLED)
                tick_count += 1
                if tick_count < max_ticks:
                    og.Controller.set(node.get_attribute("outputs:tick"), e_state.ENABLED)
                else:
                    # Finish latent naturally.
                    og.Controller.set(node.get_attribute("outputs:execOut"), e_state.LATENT_FINISH)
                    tick_count = 0

                if boop_val == e_state.ENABLED:
                    # We get here during latent ticking, if the boop input is enabled.
                    self.assertEqual(exec_in_val, e_state.DISABLED)
                    self.assertEqual(cancel_val, e_state.DISABLED)
                    boop_count += 1
                return True

            @staticmethod
            def get_node_type() -> str:
                """Get node type"""
                return "omni.graph.test.CancelTickerPy"

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                """Initialize node attributes"""
                node_type.add_input(
                    "inputs:execIn",
                    "execution",
                    True,
                )
                node_type.add_input(
                    "inputs:cancel",
                    "execution",
                    True,
                )
                node_type.add_input(
                    "inputs:boop",
                    "execution",
                    True,
                )
                node_type.add_output("outputs:tick", "execution", True)
                node_type.add_output("outputs:canceled", "execution", True)
                node_type.add_output("outputs:execOut", "execution", True)
                return True

        og.register_node_type(CancelTickerPy, 1)

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (ticker, start, _, cancel, boop, _, _, counter), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Ticker", "omni.graph.test.CancelTickerPy"),
                    ("Start", "omni.graph.action.OnImpulseEvent"),
                    ("Start2", "omni.graph.action.OnImpulseEvent"),
                    ("Cancel", "omni.graph.action.OnImpulseEvent"),
                    ("Boop", "omni.graph.action.OnImpulseEvent"),
                    ("Once", "omni.graph.action.Once"),
                    ("Once2", "omni.graph.action.Once"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("Start.outputs:execOut", "Ticker.inputs:execIn"),
                    ("Start2.outputs:execOut", "Ticker.inputs:execIn"),
                    ("Cancel.outputs:execOut", "Once.inputs:execIn"),
                    ("Once.outputs:once", "Ticker.inputs:cancel"),
                    ("Once.outputs:after", "Ticker.inputs:cancel"),
                    ("Cancel.outputs:execOut", "Once2.inputs:execIn"),
                    ("Boop.outputs:execOut", "Ticker.inputs:boop"),
                    ("Once2.outputs:once", "Ticker.inputs:cancel"),
                    ("Once2.outputs:after", "Ticker.inputs:cancel"),
                    ("Ticker.outputs:tick", "Counter.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("Start.inputs:onlyPlayback", False),
                    ("Start2.inputs:onlyPlayback", False),
                    ("Cancel.inputs:onlyPlayback", False),
                    ("Boop.inputs:onlyPlayback", False),
                ],
            },
        )
        # Cancel, check nothing happens.
        og.Controller.set(controller.attribute("state:enableImpulse", cancel), True)
        await controller.evaluate(graph)
        exec_out = og.Controller.get(controller.attribute("outputs:tick", ticker))
        self.assertEqual(exec_out, e_state.DISABLED)

        # Start ticking.
        og.Controller.set(controller.attribute("state:enableImpulse", start), True)

        await controller.evaluate(graph)  # Starts latent state.

        await controller.evaluate(graph)  # Tick 1
        self.assertEqual(tick_count, 1)
        # Verify the tick has started.
        exec_out = og.Controller.get(controller.attribute("outputs:tick", ticker))
        self.assertEqual(exec_out, e_state.ENABLED)
        await controller.evaluate(graph)  # Tick 2
        self.assertEqual(tick_count, 2)

        exec_out = og.Controller.get(controller.attribute("outputs:tick", ticker))
        self.assertEqual(exec_out, e_state.ENABLED)
        await controller.evaluate(graph)  # Tick 3
        self.assertEqual(tick_count, 3)

        # Boop - node keeps ticking.
        og.Controller.set(controller.attribute("state:enableImpulse", boop), True)
        # Boop will trigger a compute, which increments boop + ticks AND the normal latent tick.
        await controller.evaluate(graph)
        self.assertEqual(boop_count, 1)
        self.assertEqual(tick_count, 5)

        # Now check that the next tick can run WITHOUT inputs:boop being high.
        await controller.evaluate(graph)
        self.assertEqual(boop_count, 1)  # No change in boop count (OM-64856).
        self.assertEqual(tick_count, 6)

        # Now check that we can't re-trigger execIn.
        self.assertEqual(exec_in_latent_count, 0)
        og.Controller.set(controller.attribute("state:enableImpulse", start), True)
        # Start will not trigger any compute because the node is latent.
        await controller.evaluate(graph)
        self.assertEqual(exec_in_latent_count, 0)
        self.assertEqual(boop_count, 1)
        self.assertEqual(tick_count, 7)
        # Unset the impulse.
        og.Controller.set(controller.attribute("state:enableImpulse", start), False)

        # Now check the normal tick proceeds as normal.
        await controller.evaluate(graph)
        self.assertEqual(boop_count, 1)
        self.assertEqual(tick_count, 8)

        # Cancel.
        counter_attr = controller.attribute("outputs:count", counter)
        count_0 = og.Controller.get(counter_attr)
        og.Controller.set(controller.attribute("state:enableImpulse", cancel), True)
        await controller.evaluate(graph)  # Latent finish.
        await controller.evaluate(graph)  # No action.
        await controller.evaluate(graph)  # No action.
        count_1 = og.Controller.get(counter_attr)
        self.assertEqual(count_0 + 1, count_1)

    # ----------------------------------------------------------------------
    async def test_simple_expression(self):
        """Test ActionGraph simple expression of add nodes"""
        keys = og.Controller.Keys
        controller = og.Controller()
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Const", "omni.graph.nodes.ConstantDouble3"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("Add2", "omni.graph.nodes.Add"),
                    ("Add3", "omni.graph.nodes.Add"),
                    ("PostAdd", "omni.graph.nodes.Add"),
                    ("Write1", "omni.graph.nodes.WritePrimAttribute"),
                    ("Write2", "omni.graph.nodes.WritePrimAttribute"),
                    ("Write3", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: (
                    "/World/TestPrim",
                    {
                        "val1": ("double[3]", Gf.Vec3d(1, 1, 1)),
                        "val2": ("double[3]", Gf.Vec3d(2, 2, 2)),
                        "val3": ("double[3]", Gf.Vec3d(2, 2, 2)),
                    },
                ),
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Const.inputs:value", [1, 2, 3]),
                    ("Write1.inputs:name", "val1"),
                    ("Write1.inputs:primPath", "/World/TestPrim"),
                    ("Write1.inputs:usePath", True),
                    ("Write2.inputs:name", "val2"),
                    ("Write2.inputs:primPath", "/World/TestPrim"),
                    ("Write2.inputs:usePath", True),
                    ("Write3.inputs:name", "val3"),
                    ("Write3.inputs:primPath", "/World/TestPrim"),
                    ("Write3.inputs:usePath", True),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Write1.inputs:execIn"),
                    ("Write1.outputs:execOut", "Write2.inputs:execIn"),
                    ("Write2.outputs:execOut", "Write3.inputs:execIn"),
                    ("Const.inputs:value", "Add.inputs:a"),
                    ("Const.inputs:value", "Add.inputs:b"),
                    ("Const.inputs:value", "Add2.inputs:a"),
                    ("Const.inputs:value", "Add2.inputs:b"),
                    ("Add.outputs:sum", "Add3.inputs:a"),
                    ("Add2.outputs:sum", "Add3.inputs:b"),
                    ("Add3.outputs:sum", "Write1.inputs:value"),
                    ("Add3.outputs:sum", "Write2.inputs:value"),
                    ("Add3.outputs:sum", "PostAdd.inputs:a"),
                    ("Add3.outputs:sum", "PostAdd.inputs:b"),
                    ("PostAdd.outputs:sum", "Write3.inputs:value"),
                ],
            },
        )
        await controller.evaluate(graph)
        context = omni.usd.get_context()
        stage = context.get_stage()
        await omni.kit.app.get_app().next_update_async()
        self.assertListEqual([4, 8, 12], list(stage.GetAttributeAtPath("/World/TestPrim.val1").Get()))
        self.assertListEqual([4, 8, 12], list(stage.GetAttributeAtPath("/World/TestPrim.val2").Get()))
        self.assertListEqual([8, 16, 24], list(stage.GetAttributeAtPath("/World/TestPrim.val3").Get()))

    # ----------------------------------------------------------------------
    async def test_stateful_flowcontrol_evaluation(self):
        """Test that stateful flow control nodes are fully evaluated"""
        #                                 b
        #                    +----------+   +---------+
        #               +--->| Sequence +-->|Counter1 |
        #               |    +----------+   +---------+
        # +-----------+ |
        # | OnImpulse +-+
        # +-----------+ |
        #               |    +----------+   +----------+
        #               +--->| ForLoop1 +-->| Counter2 |
        #                    +----------+   +----------+
        #                              finished
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, _, counter1_node, _, counter2_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Sequence", "omni.graph.action.Sequence"),
                    ("Counter1", "omni.graph.action.Counter"),
                    ("ForLoop1", "omni.graph.action.ForLoop"),
                    ("Counter2", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "Sequence.inputs:execIn"),
                    ("Sequence.outputs:b", "Counter1.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "ForLoop1.inputs:execIn"),
                    ("ForLoop1.outputs:finished", "Counter2.inputs:execIn"),
                ],
                keys.SET_VALUES: [("OnImpulse.inputs:onlyPlayback", False), ("ForLoop1.inputs:stop", 10)],
            },
        )
        await controller.evaluate(graph)

        # Trigger graph once.
        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: ("OnImpulse.state:enableImpulse", True)})
        await controller.evaluate(graph)
        # Verify that counter was called in spite of sequence 'a' being disconnected.
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", counter1_node)), 1)
        # Verify that counter was called in spite of there being no loopBody - execution evaluator has to still trigger
        # the loop 11 times despite there being no downstream connection.
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", counter2_node)), 1)

    # ----------------------------------------------------------------------
    async def test_unresolve_on_disconnect(self):
        """Tests unresolving attribs when action node is disconnected"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, for_each_node, _, _), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("ForEach", "omni.graph.action.ForEach"),
                    ("IntArrayPrim", "omni.graph.nodes.ReadPrimAttribute"),
                    ("IntSinkPrim", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: [
                    ("/World/IntArray", {"myintarray": ("int[]", range(10))}),
                    ("/World/IntSink", {"myint": ("int", 0)}),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "ForEach.inputs:execIn"),
                    ("IntArrayPrim.outputs:value", "ForEach.inputs:arrayIn"),
                    ("ForEach.outputs:element", "IntSinkPrim.inputs:value"),
                ],
                keys.SET_VALUES: [
                    ("IntArrayPrim.inputs:name", "myintarray"),
                    ("IntArrayPrim.inputs:primPath", "/World/IntArray"),
                    ("IntArrayPrim.inputs:usePath", True),
                    ("IntSinkPrim.inputs:name", "myint"),
                    ("IntSinkPrim.inputs:primPath", "/World/IntSink"),
                    ("IntSinkPrim.inputs:usePath", True),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        await controller.evaluate(graph)
        array_attr = controller.attribute("inputs:arrayIn", for_each_node)
        element_attr = controller.attribute("outputs:element", for_each_node)
        self.assertEqual(element_attr.get_resolved_type(), og.Type(og.BaseDataType.INT, 1, 0))

        # When we disconnect all data connections, they should unresolve even though we still
        # have the execution connection in place.
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.DISCONNECT: [
                    ("IntArrayPrim.outputs:value", "ForEach.inputs:arrayIn"),
                    ("ForEach.outputs:element", "IntSinkPrim.inputs:value"),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(array_attr.get_resolved_type(), og.Type(og.BaseDataType.UNKNOWN))
        self.assertEqual(element_attr.get_resolved_type(), og.Type(og.BaseDataType.UNKNOWN))

    async def test_actiongraph_abi(self):
        """Tests IActionGraph ABI functions"""
        i_ag = get_interface()

        g_exec_out = 0

        class State(Enum):
            ENABLED = auto()
            ENABLED_AND_PUSH = auto()
            START = auto()
            END = auto()

        class _TestActionGraphPy:
            """Helper node to test behavior"""

            class _State:
                val = None

            instance_states = {}

            @staticmethod
            def initialize(context, node):
                _TestActionGraphPy.instance_states[node] = _TestActionGraphPy._State()

            @staticmethod
            def compute(_: og.GraphContext, node: og.Node) -> bool:
                """Compute method"""
                nonlocal g_exec_out
                if not i_ag.get_latent_state():
                    self.assertTrue(i_ag.get_execution_enabled("inputs:execIn"))
                self.assertFalse(i_ag.get_execution_enabled("inputs:execUnused"))
                self.assertFalse(i_ag.get_execution_enabled("inputs:boolUnused"))

                state = _TestActionGraphPy.instance_states[node]
                if state.val == State.ENABLED_AND_PUSH:
                    # Must be re-entering
                    i_ag.set_execution_enabled("outputs:finished_from_pushed")
                    state.val = None
                    return True
                if state.val == State.START:
                    # Must be in latent state tick
                    self.assertTrue(i_ag.get_latent_state())
                    i_ag.end_latent_state()
                    i_ag.set_execution_enabled("outputs:finished_from_latent")
                    state.val = None
                    return True
                if g_exec_out == State.ENABLED:
                    i_ag.set_execution_enabled("outputs:execOut")
                elif g_exec_out == State.ENABLED_AND_PUSH:
                    i_ag.set_execution_enabled_and_pushed("outputs:execOut")
                elif g_exec_out == State.START:
                    self.assertFalse(i_ag.get_latent_state())
                    i_ag.start_latent_state()
                elif g_exec_out == State.END:
                    i_ag.end_latent_state()
                state.val = g_exec_out
                return True

            @staticmethod
            def get_node_type() -> str:
                """Get node type"""
                return "omni.graph.action._TestActionGraphPy"

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                """Initialize node attributes"""
                node_type.add_input("inputs:execIn", "execution", True)
                node_type.add_input("inputs:execUnused", "execution", True)
                node_type.add_input("inputs:boolUnused", "bool", True)
                node_type.add_output("outputs:execOut", "execution", True)
                node_type.add_output("outputs:finished_from_latent", "execution", True)
                node_type.add_output("outputs:finished_from_pushed", "execution", True)

        og.register_node_type(_TestActionGraphPy, 1)

        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, _, _, _, _), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Test", "omni.graph.action._TestActionGraphPy"),
                    ("CounterOut", "omni.graph.action.Counter"),
                    ("CounterLatent", "omni.graph.action.Counter"),
                    ("CounterPushed", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Test.inputs:execIn"),
                    ("Test.outputs:execOut", "CounterOut.inputs:execIn"),
                    ("Test.outputs:finished_from_latent", "CounterLatent.inputs:execIn"),
                    ("Test.outputs:finished_from_pushed", "CounterPushed.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        await controller.evaluate(graph)

        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterOut.outputs:count"), 0)
        g_exec_out = State.ENABLED
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterOut.outputs:count"), 1)
        g_exec_out = State.ENABLED_AND_PUSH
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterOut.outputs:count"), 2)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterPushed.outputs:count"), 1)
        g_exec_out = State.START
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterOut.outputs:count"), 2)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterPushed.outputs:count"), 1)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterLatent.outputs:count"), 0)
        g_exec_out = State.END
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterOut.outputs:count"), 2)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterPushed.outputs:count"), 1)
        self.assertEqual(og.Controller.get(f"{self.TEST_GRAPH_PATH}/CounterLatent.outputs:count"), 1)
        with ogts.ExpectedError():
            with self.assertRaises(RuntimeError):
                i_ag.start_latent_state()

    # ------------------------------------------------------------------------
    async def test_node_self_destruction(self):
        """Test that we don't crash if a very bad node causes a resync of the stage"""
        # Executing this node may cause errors and fail the test as the underlying
        # authoring node is destroyed while the task is executing. We just verify that only
        # the expected error is generated.
        class _TestBadNodePy:
            @staticmethod
            def compute(_: og.GraphContext, node: og.Node) -> bool:
                context = omni.usd.get_context()
                stage = context.get_stage()
                stage.RemovePrim(node.get_prim_path())

            @staticmethod
            def get_node_type() -> str:
                """Get node type"""
                return "omni.graph.action._TestBadNodePy"

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                """Initialize node attributes"""
                node_type.add_input("inputs:execIn", "execution", True)
                node_type.add_output("outputs:execOut", "execution", True)

        og.register_node_type(_TestBadNodePy, 1)

        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Test", "omni.graph.action._TestBadNodePy"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Test.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        with ogts.ExpectedError():  # Authoring node for Test was lost
            await controller.evaluate(graph)
