"""Action Graph Evaluation Tests Part 2"""
import json

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.usd


# ======================================================================
class TestActionGraphEvaluation2(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_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):
                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)

                want_debug and print(
                    f"### {tick_count} execIn={exec_in_val} cancel={cancel_val} boop={boop_val}"
                )  # noqa: expression-not-assigned

                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:
                return "omni.graph.test.CancelTickerPy"

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

        # 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_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, (_, const_a, _, switch_a, _), _, _) = 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"),
                    ("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"]["LastNonEmptyEvaluation"]["Trace"]

        # Verify the evaluation trace includes exactly what we expect
        expected_trace = [
            const_a.get_prim_path(),
            switch_a.get_prim_path(),
        ]
        self.assertListEqual(expected_trace, trace)

    # ----------------------------------------------------------------------
    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_fan_out_shared_deps2(self):
        """Test that dependency sort works when there is shared data in exec fan-out"""
        #  ┌───────┐            ┌────────┐
        #  │Const1 ├───────────►│Append1 │
        #  └───────┘            │        ├──────────►┌───────────┐
        #                   ┌──►└────────┘           │ WriteVar1 │
        # ┌─────────────┐   │                  ┌────►└───────────┘
        # │ GraphTarget ├───┤                  │
        # └─────────────┘   └─►┌─────────┐     │     ┌───────────┐
        #                      │ Append2 ├─────┼────►│ WriteVar2 │
        # ┌────────┐        ┌─►└─────────┘     │     └───────────┘
        # │ Const2 ├────────┘                  │      ▲
        # └────────┘                           │      │
        #                                      │      │
        #                      ┌──────────┐    │      │
        #                      │ OnImpulse├────┴──────┘
        #                      └──────────┘
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Const1", "omni.graph.nodes.ConstantToken"),
                    ("Const2", "omni.graph.nodes.ConstantToken"),
                    ("Append1", "omni.graph.nodes.AppendPath"),
                    ("Append2", "omni.graph.nodes.AppendPath"),
                    ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                    ("WriteVar1", "omni.graph.core.WriteVariable"),
                    ("WriteVar2", "omni.graph.core.WriteVariable"),
                ],
                keys.CREATE_VARIABLES: [
                    ("path1", og.Type(og.BaseDataType.TOKEN)),
                    ("path2", og.Type(og.BaseDataType.TOKEN)),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "WriteVar1.inputs:execIn"),
                    ("OnImpulse.outputs:execOut", "WriteVar2.inputs:execIn"),
                    ("GraphTarget.outputs:primPath", "Append1.inputs:path"),
                    ("GraphTarget.outputs:primPath", "Append2.inputs:path"),
                    ("Const1.inputs:value", "Append1.inputs:suffix"),
                    ("Const2.inputs:value", "Append2.inputs:suffix"),
                    ("Append1.outputs:path", "WriteVar1.inputs:value"),
                    ("Append2.outputs:path", "WriteVar2.inputs:value"),
                ],
                keys.SET_VALUES: [
                    ("WriteVar1.inputs:variableName", "path1"),
                    ("WriteVar2.inputs:variableName", "path2"),
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("OnImpulse.state:enableImpulse", True),
                    ("Const1.inputs:value", "A"),
                    ("Const2.inputs:value", "B"),
                ],
            },
        )
        await controller.evaluate(graph)
        context = graph.get_default_graph_context()
        graph_path = self.TEST_GRAPH_PATH
        variable = graph.find_variable("path1")
        self.assertEquals(variable.get(context), f"{graph_path}/A")
        variable = graph.find_variable("path2")
        self.assertEquals(variable.get(context), f"{graph_path}/B")
