"""Action Graph Evaluation Tests"""
import asyncio

import carb.events
import omni.client
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd


# ======================================================================
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_exec_fan_out(self):
        """Test that fanning out from an exec port works"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("FF1", "omni.graph.action.FlipFlop"),
                    ("FF2", "omni.graph.action.FlipFlop"),
                    ("FF11", "omni.graph.action.FlipFlop"),
                    ("FF12", "omni.graph.action.FlipFlop"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "FF1.inputs:execIn"),
                    ("OnTick.outputs:tick", "FF2.inputs:execIn"),
                    ("FF1.outputs:a", "FF11.inputs:execIn"),
                    ("FF1.outputs:a", "FF12.inputs:execIn"),
                ],
            },
        )

        # 1. OnTick triggers FF1 which triggers FF11 and FF12, then FF2
        # 2. OnTick triggers FF1 and FF2
        # 3. OnTick triggers FF1 which triggers FF11 and FF12, then FF2

        await controller.evaluate(graph)
        flip_flops = nodes[1:]
        ff_state = [og.Controller.get(controller.attribute("outputs:isA", node)) for node in flip_flops]
        self.assertEqual(ff_state, [True, True, True, True])
        await controller.evaluate(graph)
        ff_state = [og.Controller.get(controller.attribute("outputs:isA", node)) for node in flip_flops]
        self.assertEqual(ff_state, [False, False, True, True])
        await controller.evaluate(graph)
        ff_state = [og.Controller.get(controller.attribute("outputs:isA", node)) for node in flip_flops]
        self.assertEqual(ff_state, [True, True, False, False])

    # ----------------------------------------------------------------------
    async def test_chained_stateful_nodes(self):
        """Test that chaining loop nodes works"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, _, _, counter_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("ForLoop1", "omni.graph.action.ForLoop"),
                    ("ForLoop2", "omni.graph.action.ForLoop"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "ForLoop1.inputs:execIn"),
                    ("ForLoop1.outputs:loopBody", "ForLoop2.inputs:execIn"),
                    ("ForLoop2.outputs:loopBody", "Counter.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("ForLoop1.inputs:stop", 5),
                    ("ForLoop2.inputs:stop", 5),
                ],
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", counter_node)), 5 * 5)

    # ----------------------------------------------------------------------
    async def test_async_nodes(self):
        """Test asynchronous action nodes"""
        # Check that a nested loop state is maintained when executing a latent delay
        #
        # +---------+   +----------+    +----------+    +-------+    +--------+
        # | IMPULSE +-->| FOR-LOOP +--->| FOR-LOOP +--->| DELAY +--->| COUNTER|
        # +---------+   +----------+    +----------+    +-------+    +--------+
        #
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, _, _, _, counter_node, _, _), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("ForLoop1", "omni.graph.action.ForLoop"),
                    ("ForLoop2", "omni.graph.action.ForLoop"),
                    ("Delay", "omni.graph.action.Delay"),
                    ("Counter", "omni.graph.action.Counter"),
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Counter2", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "ForLoop1.inputs:execIn"),
                    ("ForLoop1.outputs:loopBody", "ForLoop2.inputs:execIn"),
                    ("ForLoop2.outputs:loopBody", "Delay.inputs:execIn"),
                    ("Delay.outputs:finished", "Counter.inputs:execIn"),
                    ("OnTick.outputs:tick", "Counter2.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("Delay.inputs:duration", 0.1),
                    ("ForLoop1.inputs:stop", 2),
                    ("ForLoop2.inputs:stop", 5),
                ],
            },
        )
        await controller.evaluate(graph)

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

        # in delay now, no count
        counter_controller = og.Controller(og.Controller.attribute("outputs:count", counter_node))
        self.assertEqual(counter_controller.get(), 0)

        # wait to ensure the first 5 delays compute
        for _ in range(5):
            await asyncio.sleep(0.2)
            await controller.evaluate(graph)
        count_val = counter_controller.get()
        self.assertGreater(count_val, 4)

        # wait and verify the remainder go through
        for _ in range(5):
            await asyncio.sleep(0.1)
            await controller.evaluate(graph)
        self.assertEqual(counter_controller.get(), 10)

    # ----------------------------------------------------------------------
    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_request_driven_node(self):
        """Test that RequestDriven nodes are computed as expected"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, counter_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                keys.SET_VALUES: [("OnImpulse.inputs:onlyPlayback", False)],
                keys.CONNECT: ("OnImpulse.outputs:execOut", "Counter.inputs:execIn"),
            },
        )
        # After several updates, there should have been no compute calls
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        counter_controller = og.Controller(og.Controller.attribute("outputs:count", counter_node))
        self.assertEqual(counter_controller.get(), 0)

        # change OnImpulse state attrib. The node should now request compute
        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: ("OnImpulse.state:enableImpulse", True)})
        await controller.evaluate(graph)
        self.assertEqual(counter_controller.get(), 1)
        # more updates should not result in more computes
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        self.assertEqual(counter_controller.get(), 1)

    # ----------------------------------------------------------------------
    async def test_fan_in_exec(self):
        """Test that execution fan-in is handled correctly."""
        # The evaluator has to consider the case that gate.enter will have contradicting upstream values.
        # Gate needs to know which input is active, it needs the value of enter to be ENABLED when it
        # is triggered by OnTick, even though OnTickDisabled has set it's output to the same attrib as DISABLED.
        #
        #    +--------------+
        #    |OnImpulseEvent+---+   +-----------+
        #    +--------------+   |   |Gate       |
        #                       +--->toggle     |
        #    +--------------+       |           |
        #    |OnTick        +------>|enter      |
        #    +--------------+       +^-----exit-+
        #                            |
        #    +--------------+        |
        #    |OnTickDisabled+--------+
        #    +--------------+
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, _, gate, on_impulse_event), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("OnTickDisabled", "omni.graph.action.OnTick"),
                    ("Gate", "omni.graph.action.Gate"),
                    ("OnImpulseEvent", "omni.graph.action.OnImpulseEvent"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Gate.inputs:enter"),
                    ("OnTickDisabled.outputs:tick", "Gate.inputs:enter"),
                    ("OnImpulseEvent.outputs:execOut", "Gate.inputs:toggle"),
                ],
                keys.SET_VALUES: [
                    ("Gate.inputs:startClosed", True),
                    ("OnTick.inputs:onlyPlayback", False),
                    ("OnImpulseEvent.inputs:onlyPlayback", False),
                ],
            },
        )
        await controller.evaluate(graph)

        gate_exit = controller.attribute("outputs:exit", gate)

        # Verify the Gate has not triggered
        self.assertFalse(gate_exit.get())
        await controller.evaluate(graph)
        self.assertFalse(gate_exit.get())

        # toggle the gate and verify that the tick goes through. The first evaluate it isn't known if the Gate will
        # trigger because the order that entry points are executed is not defined... FIXME
        controller.attribute("state:enableImpulse", on_impulse_event).set(True)
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        self.assertTrue(gate_exit.get())

    # ----------------------------------------------------------------------
    async def test_fan_out_exec(self):
        """Test that execution fan-out is handled correctly."""
        # We want to reset the execution attribute states before the node compute() to avoid bugs
        # that arise when authors forget to fully specify the output states. However we can't
        # do this in the middle of traversal, because fan-out from a connection requires that the state
        # be preserved for every downstream node which may read from it (like Gate).

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, gate, _), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Gate", "omni.graph.action.Gate"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Counter.inputs:execIn"),
                    ("OnTick.outputs:tick", "Gate.inputs:enter"),
                ],
                keys.SET_VALUES: [
                    ("Gate.inputs:startClosed", False),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        await controller.evaluate(graph)

        gate_exit = controller.attribute("outputs:exit", gate)

        # Verify the Gate has triggered
        self.assertTrue(gate_exit.get())
        await controller.evaluate(graph)
        self.assertTrue(gate_exit.get())

    # ----------------------------------------------------------------------
    async def test_onclosing(self):
        """Test OnClosing node"""

        # Test OnClosing is tricky because OG is being destroyed when it happens -
        # so test by sending a custom event when the network is triggered
        # and then checking if we got that event

        def registered_event_name(event_name):
            """Returns the internal name used for the given custom event name"""
            n = "omni.graph.action." + event_name
            return carb.events.type_from_string(n)

        got_event = [False]

        def on_event(_):
            got_event[0] = True

        reg_event_name = registered_event_name("foo")
        message_bus = omni.kit.app.get_app().get_message_bus_event_stream()
        sub = message_bus.create_subscription_to_push_by_type(reg_event_name, on_event)
        self.assertIsNotNone(sub)

        controller = og.Controller()
        keys = og.Controller.Keys
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnClosing", "omni.graph.action.OnClosing"),
                    ("Send", "omni.graph.action.SendCustomEvent"),
                ],
                keys.CONNECT: [("OnClosing.outputs:execOut", "Send.inputs:execIn")],
                keys.SET_VALUES: [("Send.inputs:eventName", "foo"), ("Send.inputs:path", "Test Path")],
            },
        )
        # evaluate once so that graph is in steady state
        await controller.evaluate()
        # close the stage
        usd_context = omni.usd.get_context()
        (result, _) = await usd_context.close_stage_async()
        self.assertTrue(result)
        # Check our handler was called
        self.assertTrue(got_event[0])

    async def test_onloaded(self):
        """Test OnLoaded node"""

        def registered_event_name(event_name):
            """Returns the internal name used for the given custom event name"""
            n = "omni.graph.action." + event_name
            return carb.events.type_from_string(n)

        events = []

        def on_event(e):
            events.append(e.payload["!path"])

        reg_event_name = registered_event_name("foo")
        message_bus = omni.kit.app.get_app().get_message_bus_event_stream()
        sub = message_bus.create_subscription_to_push_by_type(reg_event_name, on_event)
        self.assertIsNotNone(sub)

        controller = og.Controller()
        keys = og.Controller.Keys
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("OnLoaded", "omni.graph.action.OnLoaded"),
                    ("Send1", "omni.graph.action.SendCustomEvent"),
                    ("Send2", "omni.graph.action.SendCustomEvent"),
                ],
                keys.CONNECT: [
                    ("OnLoaded.outputs:execOut", "Send1.inputs:execIn"),
                    ("OnTick.outputs:tick", "Send2.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Send1.inputs:eventName", "foo"),
                    ("Send2.inputs:eventName", "foo"),
                    ("Send1.inputs:path", "Loaded"),
                    ("Send2.inputs:path", "Tick"),
                ],
            },
        )
        # evaluate once so that graph is in steady state
        await controller.evaluate()

        # Verify Loaded came before OnTick
        self.assertListEqual(events, ["Loaded", "Tick"])

    # ----------------------------------------------------------------------
    async def test_active_latent(self):
        """exercise a latent node that executes downstream nodes while latent"""

        # +--------+   +----------+finished+-------------+
        # |  OnTick+-->| TickN    +-------->FinishCounter|
        # +--------+   |          |        +-------------+
        #              |          +-+
        #              +----------+ |     +------------+     +------------+    +------------+
        #                           +-----> TickCounter+----->TickCounter2+---->TickCounter3|
        #                          tick   +------------+     +------------+    +------------+

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

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("TickN", "omni.graph.action.TickN"),
                    ("FinishCounter", "omni.graph.action.Counter"),
                    ("TickCounter", "omni.graph.action.Counter"),
                    ("TickCounter2", "omni.graph.action.Counter"),
                    ("TickCounter3", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "TickN.inputs:execIn"),
                    ("TickN.outputs:finished", "FinishCounter.inputs:execIn"),
                    ("TickN.outputs:tick", "TickCounter.inputs:execIn"),
                    ("TickCounter.outputs:execOut", "TickCounter2.inputs:execIn"),
                    ("TickCounter2.outputs:execOut", "TickCounter3.inputs:execIn"),
                ],
                keys.SET_VALUES: [("TickN.inputs:duration", 3), ("OnTick.inputs:onlyPlayback", False)],
            },
        )

        (_, _, finish_counter, tick_counter, _, tick_counter_3) = nodes
        finish_counter_controller = og.Controller(og.Controller.attribute("outputs:count", finish_counter))
        tick_counter_controller = og.Controller(og.Controller.attribute("outputs:count", tick_counter))
        tick_counter_3_controller = og.Controller(og.Controller.attribute("outputs:count", tick_counter_3))

        await controller.evaluate(graph)
        self.assertEqual(finish_counter_controller.get(), 0)
        await controller.evaluate(graph)
        self.assertEqual(tick_counter_controller.get(), 1)
        await controller.evaluate(graph)
        self.assertEqual(finish_counter_controller.get(), 0)
        self.assertEqual(tick_counter_controller.get(), 2)
        await controller.evaluate(graph)
        self.assertEqual(finish_counter_controller.get(), 0)
        self.assertEqual(tick_counter_3_controller.get(), 3)
        await controller.evaluate(graph)
        self.assertEqual(finish_counter_controller.get(), 1)
        self.assertEqual(tick_counter_3_controller.get(), 3)
        await controller.evaluate(graph)
        self.assertEqual(finish_counter_controller.get(), 1)
        self.assertEqual(tick_counter_3_controller.get(), 3)

    # ----------------------------------------------------------------------
    async def test_latent_chain(self):
        """exercise a chain of latent nodes"""

        # +---------+  +-------+  tick  +-------+   tick  +-------+
        # |OnImpulse+-->TickA  +-------->TickB  +-------->|LatentC|
        # +---------+  +-----+-+        +------++         +-------+-----+
        #                    |          finish |                finish  |
        #            finish  | +-------------+ |   +-------------+    +-v----------+
        #                    +->TickCounterA | +-->| TickCounterB|    |TickCounterC|
        #                      +-------------+     +-------------+    +------------+

        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.TickN"),
                    ("TickB", "omni.graph.action.TickN"),
                    ("LatentC", "omni.graph.action.TickN"),
                    ("TickCounterA", "omni.graph.action.Counter"),
                    ("TickCounterB", "omni.graph.action.Counter"),
                    ("TickCounterC", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "TickA.inputs:execIn"),
                    ("TickA.outputs:tick", "TickB.inputs:execIn"),
                    ("TickA.outputs:finished", "TickCounterA.inputs:execIn"),
                    ("TickB.outputs:tick", "LatentC.inputs:execIn"),
                    ("TickB.outputs:finished", "TickCounterB.inputs:execIn"),
                    ("LatentC.outputs:finished", "TickCounterC.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("TickA.inputs:duration", 2),
                    ("TickB.inputs:duration", 2),
                    ("LatentC.inputs:duration", 2),
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("OnImpulse.state:enableImpulse", True),
                ],
            },
        )

        (_, _, _, _, tick_counter_a, tick_counter_b, tick_counter_c) = nodes

        for _ in range(16):
            await controller.evaluate(graph)

        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", tick_counter_a)), 1)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", tick_counter_b)), 2)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", tick_counter_c)), 4)

    # ----------------------------------------------------------------------
    async def test_latent_and_push(self):
        """exercise latent nodes in combination with stateful loop node"""
        #
        # +---------+   +-------+     tick  +--------+  loopBody +-------+    +------------+
        # |OnImpulse+-->|TickA  +----------->ForLoop1++--------->|TickB  +-+->|TickCounterB|
        # +---------+   +----+--+           +--------++          +-------+ |  +------------+
        #                    | finish                 |                    |
        #                    |                        |                    |
        #                    |      +--------------+ +v----------------+ +-v------------+
        #                    +----->|FinishCounterA| |FinishLoopCounter| |FinishCounterB|
        #                           +--------------+ +-----------------+ +--------------+

        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.TickN"),
                    ("ForLoop1", "omni.graph.action.ForLoop"),
                    ("TickB", "omni.graph.action.TickN"),
                    ("FinishLoopCounter", "omni.graph.action.Counter"),
                    ("FinishCounterA", "omni.graph.action.Counter"),
                    ("FinishCounterB", "omni.graph.action.Counter"),
                    ("TickCounterB", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "TickA.inputs:execIn"),
                    ("TickA.outputs:tick", "ForLoop1.inputs:execIn"),
                    ("TickA.outputs:finished", "FinishCounterA.inputs:execIn"),
                    ("ForLoop1.outputs:loopBody", "TickB.inputs:execIn"),
                    ("ForLoop1.outputs:finished", "FinishLoopCounter.inputs:execIn"),
                    ("TickB.outputs:tick", "TickCounterB.inputs:execIn"),
                    ("TickB.outputs:finished", "FinishCounterB.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("ForLoop1.inputs:start", 0),
                    ("ForLoop1.inputs:stop", 3),
                    ("TickA.inputs:duration", 2),
                    ("TickB.inputs:duration", 2),
                    ("OnImpulse.state:enableImpulse", True),
                    ("OnImpulse.inputs:onlyPlayback", False),
                ],
            },
        )

        (_, _, _, _, finish_loop_counter, finish_counter_a, finish_counter_b, tick_counter_b) = nodes

        for _ in range(20):
            await controller.evaluate(graph)

        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", finish_counter_a)), 1)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", tick_counter_b)), 12)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", finish_counter_b)), 6)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", finish_loop_counter)), 2)

    # ----------------------------------------------------------------------
    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.TickN"),
                    ("TickB", "omni.graph.action.TickN"),
                    ("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(c, a, b, f_a, f_b):
            for node, expected in (
                (tick_counter, c),
                (tick_counter_a, a),
                (tick_counter_b, 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_diamond_fan(self):
        """Test latent nodes in parallel fan-out which fan-in to a downstream node"""
        #                                +--------++    +----------+
        #                             +--> TickA   +--->|FinishedA |---+
        #                             |  +---------+    +----------+   |
        # +---------+   +-----------+ |                                |   +------------+
        # |OnImpulse+-->|TickCounter+-+                                +-->|MergeCounter|
        # +---------+   +-----------+ |                                |   +------------+
        #                             |   +---------+    +----------+  |
        #                             +-->| TickB   +--->|FinishedB |--+
        #                                 +--------++    +----------+
        #                             |   +---------+
        #                             +-->| TickC   |
        #                                 +--------++
        controller = og.Controller()
        keys = og.Controller.Keys

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

        (_, tick_counter, _, _, tick_c, finish_counter_a, finish_counter_b, merge_counter) = nodes

        def check_counts(tc, f_a, f_b, mc, tick_c_count):
            for node, expected in (
                (tick_counter, tc),
                (finish_counter_a, f_a),
                (finish_counter_b, f_b),
                (merge_counter, mc),
            ):
                count = og.Controller.get(controller.attribute("outputs:count", node))
                self.assertEqual(count, expected, node.get_prim_path())
            self.assertEqual(tick_c.get_compute_count(), tick_c_count)

        self.assertEqual(tick_c.get_compute_count(), 0)
        # set up latent tickers
        await controller.evaluate(graph)
        check_counts(1, 0, 0, 0, 1)
        # latent ticks
        await controller.evaluate(graph)
        check_counts(1, 0, 0, 0, 2)
        # both branches complete
        await controller.evaluate(graph)
        check_counts(1, 1, 1, 2, 3)
        # no count changes + no additional computes of tickC
        await controller.evaluate(graph)
        check_counts(1, 1, 1, 2, 3)

    # ----------------------------------------------------------------------
    async def test_diamond_latent_fan(self):
        """Test latent nodes in parallel fan-out which fan-in to a latent downstream node"""
        #                  +--------++
        #               +--> TickA   +--+
        #               |  +---------+  |
        # +---------+   |               |   +-------+      +-------+
        # |OnImpulse+-->+               +-->|TickD  +-+--->|CountF |
        # +---------+   |               |   +-------+ |    +-------+
        #               |   +--------+  |             +--->+-------+
        #               +-->| TickB  +--+                  |TickE  |
        #               |   +--------+                +--->+-------+
        #               |   +--------+                |
        #               +-->| TickC  +----------------+
        #                   +--------+

        # Note that when TickA triggers TickD into latent state, then TickB hits TickD subsequently. This subsequent
        # evaluation is _transient_. Meaning that TickB will not block on a new copy of TickD.
        # This is because there is only one TickD so there can be only one state (latent or not).

        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.TickN"),
                    ("TickB", "omni.graph.action.TickN"),
                    ("TickC", "omni.graph.action.TickN"),
                    ("TickD", "omni.graph.action.TickN"),
                    ("TickE", "omni.graph.action.TickN"),
                    ("CountF", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "TickA.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "TickB.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "TickC.inputs:execIn"),
                    ("TickA.outputs:finished", "TickD.inputs:execIn"),
                    ("TickB.outputs:finished", "TickD.inputs:execIn"),
                    ("TickC.outputs:finished", "TickE.inputs:execIn"),
                    ("TickD.outputs:finished", "TickE.inputs:execIn"),
                    ("TickD.outputs:finished", "CountF.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("TickA.inputs:duration", 1),
                    ("TickB.inputs:duration", 1),
                    ("TickC.inputs:duration", 2),
                    ("TickD.inputs:duration", 1),
                    ("TickE.inputs:duration", 1),
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("OnImpulse.state:enableImpulse", True),
                ],
            },
        )

        (_, tick_a, tick_b, tick_c, tick_d, tick_e, count_f) = nodes

        def check_counts(i, ta, tb, tc, td, te):
            # print(f"{[node.get_compute_count() for node, expected in ((tick_a, ta), (tick_b, tb), (tick_c, tc), (tick_d, td), (tick_e, te))]}")
            for node, expected in ((tick_a, ta), (tick_b, tb), (tick_c, tc), (tick_d, td), (tick_e, te)):
                self.assertEqual(node.get_compute_count(), expected, f"Check {i} for {node.get_prim_path()}")

        #    A, B, C, D, E
        compute_counts = [
            (1, 1, 1, 0, 0),  # 0. fan out to trigger A, B, C into latent state
            (2, 2, 2, 0, 0),  # 1. A, B, C tick
            (3, 3, 3, 2, 0),  # 2. A, B end latent, D into latent via A or B, D ticks via A or B, C ticks
            (3, 3, 4, 3, 2),  # 3.
            (3, 3, 4, 3, 3),  # 4.
            (3, 3, 4, 3, 3),  # 5.
            (3, 3, 4, 3, 3),  # 6.
        ]

        for i, cc in enumerate(compute_counts):
            await controller.evaluate(graph)
            check_counts(i, *cc)
        # Verify that CountF has computed 1x due to the fan-in at TickD NOT acting like separate threads
        self.assertEqual(count_f.get_compute_count(), 1)

    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:
            @staticmethod
            def compute(context: og.GraphContext, node: og.Node):
                nonlocal i
                go = node.get_attribute("inputs:go")
                go_val = og.Controller.get(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:
                return "omni.graph.test.OnForEachEventPy"

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                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:
            @staticmethod
            def compute(context: og.GraphContext, node: og.Node):
                og.Controller.set(node.get_attribute("outputs:execOut"), og.ExecutionAttributeState.ENABLED)
                return True

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

            @staticmethod
            def initialize_type(node_type: og.NodeType):
                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)
