"""Action Graph Evaluation Tests, Part 1"""
import asyncio
import json

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_active_latent(self):
        """Exercise a latent node that executes downstream nodes while latent"""

        # +--------+   +----------+finished+-------------+
        # |  OnTick+-->| Countdown+-------->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"),
                    ("Countdown", "omni.graph.action.Countdown"),
                    ("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", "Countdown.inputs:execIn"),
                    ("Countdown.outputs:finished", "FinishCounter.inputs:execIn"),
                    ("Countdown.outputs:tick", "TickCounter.inputs:execIn"),
                    ("TickCounter.outputs:execOut", "TickCounter2.inputs:execIn"),
                    ("TickCounter2.outputs:execOut", "TickCounter3.inputs:execIn"),
                ],
                keys.SET_VALUES: [("Countdown.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_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_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_cycle_break(self):
        """Test that an illegal cycle issues a warning"""

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

        (graph, (on_impulse, count_a, count_b), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("A", "omni.graph.action.Counter"),
                    ("B", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "A.inputs:execIn"),
                    ("A.outputs:execOut", "B.inputs:execIn"),
                    ("B.outputs:execOut", "A.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnImpulse.state:enableImpulse", True),
                    ("OnImpulse.inputs:onlyPlayback", False),
                ],
            },
        )
        with ogts.ExpectedError():
            await controller.evaluate(graph)
        og.Controller.set(controller.attribute("state:enableImpulse", on_impulse), True)
        with ogts.ExpectedError():
            await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", count_a)), 2)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:count", count_b)), 2)

    # ----------------------------------------------------------------------
    async def test_dep_sort_fan_out(self):
        """Test that dependency sort works when there is data fan-out"""
        #                      +-------------+
        #            +-------->|             |
        #            |         | SwitchTokenA|
        #            |    +--->+-------------+
        # +----------+    |
        # |OnImpulse +----|------+           +--------------+
        # +----------+    |      +---------->| SwitchTokenB |
        #                 |                 +^-------------+
        #          +------+-+   +--------+   |
        #          | ConstA +--->AppendB +---+
        #          +--------+   +--------+

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

        (graph, (_, _, _, _, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("ConstA", "omni.graph.nodes.ConstantToken"),
                    ("AppendB", "omni.graph.nodes.AppendString"),
                    ("SwitchTokenA", "omni.graph.action.SwitchToken"),
                    ("SwitchTokenB", "omni.graph.action.SwitchToken"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "SwitchTokenA.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "SwitchTokenB.inputs:execIn"),
                    ("ConstA.inputs:value", "SwitchTokenA.inputs:value"),
                    ("ConstA.inputs:value", "AppendB.inputs:value"),
                    ("AppendB.outputs:value", "SwitchTokenB.inputs:value"),
                ],
                keys.SET_VALUES: [
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("OnImpulse.state:enableImpulse", True),
                    ("AppendB.inputs:suffix", {"value": "Foo", "type": "token"}),
                ],
            },
        )
        await controller.evaluate(graph)
        graph_state = og.OmniGraphInspector().as_json(graph, flags=["evaluation"])
        graph_state_obj = json.loads(graph_state)
        trace = graph_state_obj["Evaluator"]["Instances"][0]["LastNonEmptyEvaluation"]["Trace"]
        # The switches can run in any order
        self.assertTrue(
            trace in (["SwitchTokenA", "AppendB", "SwitchTokenB"], ["AppendB", "SwitchTokenB", "SwitchTokenA"])
        )

    # ----------------------------------------------------------------------
    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.Countdown"),
                    ("TickB", "omni.graph.action.Countdown"),
                    ("TickC", "omni.graph.action.Countdown"),
                    ("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(t_c, f_a, f_b, m_c, tick_c_count):
            for node, expected in (
                (tick_counter, t_c),
                (finish_counter_a, f_a),
                (finish_counter_b, f_b),
                (merge_counter, m_c),
            ):
                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.Countdown"),
                    ("TickB", "omni.graph.action.Countdown"),
                    ("TickC", "omni.graph.action.Countdown"),
                    ("TickD", "omni.graph.action.Countdown"),
                    ("TickE", "omni.graph.action.Countdown"),
                    ("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, t_a, t_b, t_c, t_d, t_e):
            for node, expected in ((tick_a, t_a), (tick_b, t_b), (tick_c, t_c), (tick_d, t_d), (tick_e, t_e)):
                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, c_c in enumerate(compute_counts):
            await controller.evaluate(graph)
            check_counts(i, *c_c)
        # 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_dynamic_exec_pins(self):
        """Test that adding execution pins to a non-action node works"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (on_tick, to_string), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("ToString", "omni.graph.nodes.ToString"),
                ],
                keys.SET_VALUES: [
                    ("ToString.inputs:value", 42, "double"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        # Verify to_string has not been computed.
        await controller.evaluate()
        self.assertEqual(0, to_string.get_compute_count())
        self.assertEqual(1, on_tick.get_compute_count())
        # Add execution attribs and verify it still doesn't get computed.
        attrib = og.Controller.create_attribute(
            to_string,
            "inputs:execIn",
            "execution",
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
        )
        self.assertIsNotNone(attrib)
        await controller.evaluate()
        self.assertEqual(0, to_string.get_compute_count())
        self.assertEqual(2, on_tick.get_compute_count())
        # Hook up to OnTick and verify it is now computing.
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "ToString.inputs:execIn"),
                ]
            },
        )
        for i in range(10):
            await controller.evaluate()
            self.assertEqual(i + 1, to_string.get_compute_count())
            self.assertEqual(i + 3, on_tick.get_compute_count())

    # ----------------------------------------------------------------------
    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_exec_fan_out_shared_deps(self):
        """Test that dependency sort works when there is shared data in exec fan-out"""
        #                               +---------+
        #                   +---------->|  Write1 |
        #                   |           +----^----+
        #                   |                |
        #                   |     +----------+
        #                   |     |
        # +-----------+     |     |
        # | OnImpulse +-----+-----+----> +---------+
        # +-----------+     |     |      |  Write2 |
        #                   |     +----->+---------+
        #                   |     |
        #                   |     |      +---------+
        #                   +-----+----->| Write3  |
        #                         |      +---------+
        #                         |        ^
        #      +-------+      +---+----+---+
        #      | Const +----->|   Inc  |
        #      +-------+      +--------+
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Const", "omni.graph.nodes.ConstantDouble"),
                    ("Inc", "omni.graph.nodes.Increment"),
                    ("Write1", "omni.graph.nodes.WritePrimAttribute"),
                    ("Write2", "omni.graph.nodes.WritePrimAttribute"),
                    ("Write3", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: [
                    ("/World/TestPrim1", {"val": ("double", 1.0)}),
                    ("/World/TestPrim2", {"val": ("double", 2.0)}),
                    ("/World/TestPrim3", {"val": ("double", 3.0)}),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "Write1.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "Write2.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "Write3.inputs:execIn"),
                    ("Const.inputs:value", "Inc.inputs:value"),
                    ("Inc.outputs:result", "Write1.inputs:value"),
                    ("Inc.outputs:result", "Write2.inputs:value"),
                    ("Inc.outputs:result", "Write3.inputs:value"),
                ],
                keys.SET_VALUES: [
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("OnImpulse.state:enableImpulse", True),
                    ("Const.inputs:value", 41.0),
                    ("Inc.inputs:increment", 1.0),
                    ("Write1.inputs:primPath", "/World/TestPrim1"),
                    ("Write1.inputs:usePath", True),
                    ("Write1.inputs:name", "val"),
                    ("Write2.inputs:primPath", "/World/TestPrim2"),
                    ("Write2.inputs:usePath", True),
                    ("Write2.inputs:name", "val"),
                    ("Write3.inputs:primPath", "/World/TestPrim3"),
                    ("Write3.inputs:usePath", True),
                    ("Write3.inputs:name", "val"),
                ],
            },
        )
        await controller.evaluate(graph)
        stage = omni.usd.get_context().get_stage()
        for i in (1, 2, 3):
            self.assertEqual(stage.GetAttributeAtPath(f"/World/TestPrim{i}.val").Get(), 42.0)

    # ----------------------------------------------------------------------
    async def test_exec_sort_failure(self):
        """Test that sorting dependencies with non-trivial authored graph"""

        # Our global sort excludes exec nodes, so a global topo (Kahn) sort will fail such that Inc3 doesn't get
        # computed until after Add2, so instead we sort each dep network independently. This test verifies the case
        # where that matters.
        #
        #      +-----------------------------> Write1(var) +----------------------------------------+
        #      |                                ^     |                                             |
        #      |                                |     |                                             v
        # OnTick --------------------+          |     +-----------Inc------------+                Write2(var2)
        #                            |          |                                |                  ^
        #                            v          |                                |                  |
        #  Read1(var)------------> Add1 --Inc2--+                                v                  |
        #                                                  Inc3 --------------> Add2 ---------------+
        controller = og.Controller()
        keys = og.Controller.Keys

        (_, (on_tick, a_1, a_2, _, _, _, _, _, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("A1", "omni.graph.nodes.Add"),
                    ("A2", "omni.graph.nodes.Add"),
                    ("Write1", "omni.graph.core.WriteVariable"),
                    ("Write2", "omni.graph.core.WriteVariable"),
                    ("Read1", "omni.graph.core.ReadVariable"),
                    ("Inc", "omni.graph.nodes.Increment"),
                    ("Inc2", "omni.graph.nodes.Increment"),
                    ("Inc3", "omni.graph.nodes.Increment"),
                ],
                keys.CREATE_VARIABLES: [
                    ("var", og.Type(og.BaseDataType.DOUBLE)),
                    ("var2", og.Type(og.BaseDataType.DOUBLE)),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Inc3.inputs:value", {"type": "double", "value": 42.0}),
                    ("Write1.inputs:variableName", "var"),
                    ("Write2.inputs:variableName", "var2"),
                    ("Read1.inputs:variableName", "var"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Write1.inputs:execIn"),
                    ("OnTick.outputs:timeSinceStart", "A1.inputs:a"),
                    ("Read1.outputs:value", "A1.inputs:b"),
                    ("A1.outputs:sum", "Inc2.inputs:value"),
                    ("Inc2.outputs:result", "Write1.inputs:value"),
                    ("Write1.outputs:execOut", "Write2.inputs:execIn"),
                    ("Write1.outputs:value", "Inc.inputs:value"),
                    ("Inc.outputs:result", "A2.inputs:a"),
                    ("Inc3.outputs:result", "A2.inputs:b"),
                    ("A2.outputs:sum", "Write2.inputs:value"),
                ],
            },
        )
        await omni.kit.app.get_app().next_update_async()

        a_1_v = og.Controller.get(controller.attribute("outputs:sum", a_1))
        a_2_v = og.Controller.get(controller.attribute("outputs:sum", a_2))
        on_tick_dt = og.Controller.get(controller.attribute("outputs:timeSinceStart", on_tick))

        a_1_expected = 0 + on_tick_dt
        a_2_expected = (a_1_expected + 1.0 + 1.0) + (42.0 + 1.0)
        self.assertAlmostEqual(a_1_v, a_1_expected, places=3)
        self.assertAlmostEqual(a_2_v, a_2_expected, places=3)

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

    async def test_fan_in(self):
        """Test that fan-in of execution connections works as expected (from a loaded test .usda file)"""

        (result, error) = await ogts.load_test_file("TestActionFanIn.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph_path = "/World/ActionGraph"
        controller = og.Controller()
        graph = controller.graph(graph_path)
        # Trigger the loop.
        og.Controller.set(controller.attribute(f"{graph_path}/on_impulse_event.state:enableImpulse"), True)
        await controller.evaluate(graph)
        graph_state = og.OmniGraphInspector().as_json(controller.graph(graph_path), flags=["evaluation"])
        graph_state_obj = json.loads(graph_state)
        trace = graph_state_obj["Evaluator"]["Instances"][0]["LastNonEmptyEvaluation"]["Trace"]

        # Verify the first loop iteration.
        self.assertEqual("for_loop", trace[0])
        # These nodes can compute in any order
        self.assertEqual(["counter", "counter_01"], sorted(trace[1:3]))
        expected_trace = [
            "to_uint64",
            "sync_gate",
            "to_uint64",
            "sync_gate",
        ]
        self.assertListEqual(expected_trace, trace[3:7])
        trace[0:7] = []

        # Verify downstream from sync gate.
        expected_trace = [
            "counter_02",
        ]
        self.assertListEqual(expected_trace, trace[0:1])
        trace[0 : len(expected_trace)] = []
        # Verify second iteration.
        self.assertEqual("for_loop", trace[0])
        # These nodes can compute in any order
        self.assertEqual(["counter", "counter_01"], sorted(trace[1:3]))
        expected_trace = [
            "to_uint64",
            "sync_gate",
            "to_uint64",
            "sync_gate",
        ]
        self.assertListEqual(expected_trace, trace[3:7])

    # ----------------------------------------------------------------------
    async def test_loading_type_resol(self):
        """Test that loading a file with weird type resolution pattern works"""

        (result, error) = await ogts.load_test_file("load_with_type_resol.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph_path = "/World/ActionGraph"
        controller = og.Controller()
        graph = controller.graph(graph_path)
        # eval
        await controller.evaluate(graph)
        # check result
        var = graph.find_variable("Result")
        val = var.get_array(graph.get_default_graph_context(), False, 0)
        self.assertTrue((val == [(-50, -50, -50), (50, 50, 50)]).all())

    # ----------------------------------------------------------------------
    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_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.Countdown"),
                    ("ForLoop1", "omni.graph.action.ForLoop"),
                    ("TickB", "omni.graph.action.Countdown"),
                    ("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_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.Countdown"),
                    ("TickB", "omni.graph.action.Countdown"),
                    ("LatentC", "omni.graph.action.Countdown"),
                    ("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)
