from typing import List, Set

import omni.graph.core as og
import omni.kit.test
import omni.usd
from omni.graph.core import traverse_downstream_graph, traverse_upstream_graph


class TestTraversal(omni.graph.core.tests.OmniGraphTestCase):
    """Test cases for graph traversal functions."""

    async def test_downstream_connections_works(self):
        """
        Validate that downstream connections are collected correctly.
        """
        stage = omni.usd.get_context().get_stage()

        # Test graph set up
        #   [G1In1]-->[G1M1]-->[G1M2]--->[G1M4]--->[G1Out1]
        #   [G1In2]---^    |---[G1M3]----|    |--->[G1Out2]
        #
        #   [G2In1]-->[G2M1]-->[G2Out1]
        #
        #   [G3In1]
        controller = og.Controller()
        keys = controller.Keys
        (_, nodes, _, _) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("G1In1", "omni.graph.nodes.ConstantFloat"),
                    ("G1In2", "omni.graph.nodes.ConstantFloat"),
                    ("G1M1", "omni.graph.nodes.Add"),
                    ("G1M2", "omni.graph.nodes.Add"),
                    ("G1M3", "omni.graph.nodes.Add"),
                    ("G1M4", "omni.graph.nodes.Add"),
                    ("G1Out1", "omni.graph.nodes.Magnitude"),
                    ("G1Out2", "omni.graph.nodes.Magnitude"),
                    ("G2In1", "omni.graph.nodes.ConstantFloat"),
                    ("G2M1", "omni.graph.nodes.Magnitude"),
                    ("G2Out1", "omni.graph.nodes.Magnitude"),
                    ("G3In1", "omni.graph.nodes.ConstantFloat"),
                ],
                keys.CONNECT: [
                    ("G1In1.inputs:value", "G1M1.inputs:a"),
                    ("G1In2.inputs:value", "G1M1.inputs:b"),
                    ("G1M1.outputs:sum", "G1M2.inputs:a"),
                    ("G1M1.outputs:sum", "G1M2.inputs:b"),
                    ("G1M1.outputs:sum", "G1M3.inputs:a"),
                    ("G1M1.outputs:sum", "G1M3.inputs:b"),
                    ("G1M2.outputs:sum", "G1M4.inputs:a"),
                    ("G1M3.outputs:sum", "G1M4.inputs:b"),
                    ("G1M4.outputs:sum", "G1Out1.inputs:input"),
                    ("G1M4.outputs:sum", "G1Out2.inputs:input"),
                    ("G2In1.inputs:value", "G2M1.inputs:input"),
                    ("G2M1.outputs:magnitude", "G2Out1.inputs:input"),
                ],
            },
        )

        prims = tuple(stage.GetPrimAtPath(n.get_prim_path()) for n in nodes)
        _, _, _, _, g1m3, g1m4, g1out1, g1out2, _, g2m1, g2out1, _ = prims

        # Valdate the set of visited nodes
        expected = [g1m3, g1m4, g1out1, g1out2, g2m1, g2out1]
        visited = traverse_downstream_graph([g2m1, g1m3])
        actual = [og.Controller.prim(n) for n in visited]
        self.assertCountEqual(expected, actual, "Unexpected set of visited nodes downstream")

    async def test_downstream_connections_can_include_compound_graph_prim(self):
        """Test that calling for downstream connections with compounds can include the compound graph prim."""

        node_type = "omni.graph.nodes.Add"
        controller = og.Controller()
        keys = controller.Keys

        # Test graph set up
        #       |============================|
        # [O1]->|->[I1]---->|I2|----->|I3|-->|--->|O2]
        #       |---------->|  |             |
        #       |--->[I4]                    |
        #       |    [I5]->[I6]              |
        #       |============================|
        (_, _, _, nodes) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("Outer1", node_type),
                    ("Outer2", node_type),
                    (
                        "Compound",
                        {
                            keys.CREATE_NODES: [
                                ("Inner1", node_type),
                                ("Inner2", node_type),
                                ("Inner3", node_type),
                                ("Inner4", node_type),
                                ("Inner5", node_type),
                                ("Inner6", node_type),
                            ],
                            keys.PROMOTE_ATTRIBUTES: [
                                ("Inner1.inputs:a", "inputs:one"),
                                ("Inner2.inputs:a", "inputs:two"),
                                ("Inner3.outputs:sum", "outputs:value"),
                                ("Inner4.inputs:a", "inputs:three"),
                            ],
                            keys.CONNECT: [
                                ("Inner1.outputs:sum", "Inner2.inputs:b"),
                                ("Inner2.outputs:sum", "Inner3.inputs:a"),
                                ("Inner5.outputs:sum", "Inner6.inputs:a"),
                            ],
                        },
                    ),
                ],
                keys.CONNECT: [
                    ("Outer1.outputs:sum", "Compound.inputs:one"),
                    ("Compound.outputs:value", "Outer2.inputs:a"),
                ],
            },
        )

        def assert_valid_traversal(items: List, expected_results: Set[og.Node]):
            """Helper method to validate traversal results."""
            prims = [og.Controller.prim(i) for i in items]

            actual_set = traverse_downstream_graph(prims)
            self.assertEquals(expected_results, actual_set, "Unexpected set of visited nodes downstream")

        def assert_invalid_traversal(items):
            prims = [og.Controller.prim(i) for i in items]
            with self.assertRaises(og.OmniGraphError):
                traverse_downstream_graph(prims)

        compound_graph = nodes["Compound"].get_compound_graph_instance()

        # test traversal with compound graph
        assert_valid_traversal([nodes["Outer1"]], {nodes["Outer1"], nodes["Compound"], nodes["Outer2"]})

        # test with just the compound graph
        assert_valid_traversal([compound_graph], {nodes["Inner1"], nodes["Inner2"], nodes["Inner3"], nodes["Inner4"]})

        # test with a node inside the compound graph
        assert_valid_traversal([compound_graph, nodes["Inner5"]], set(compound_graph.get_nodes()))

        # invalid traversals
        assert_invalid_traversal([nodes["Inner1"], nodes["Outer1"]])  # mixed graphs
        assert_invalid_traversal([compound_graph, nodes["Outer1"]])  # compound and outer node
        assert_invalid_traversal(["/World"])  # invalid node
        assert_invalid_traversal(["/World/Testgraph"])  # non compound graph

    async def test_upstream_connections_works(self):
        """
        Validate that upstream connections are collected correctly.
        """
        stage = omni.usd.get_context().get_stage()

        # Test graph set up
        #   [G1In1]-->[G1M1]-->[G1M2]--->[G1M4]--->[G1Out1]
        #   [G1In2]---^    |---[G1M3]----|    |--->[G1Out2]
        #
        #   [G2In1]-->[G2M1]-->[G2Out1]
        #
        #   [G3In1]
        controller = og.Controller()
        keys = controller.Keys
        (_, nodes, _, _) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("G1In1", "omni.graph.nodes.ConstantFloat"),
                    ("G1In2", "omni.graph.nodes.ConstantFloat"),
                    ("G1M1", "omni.graph.nodes.Add"),
                    ("G1M2", "omni.graph.nodes.Add"),
                    ("G1M3", "omni.graph.nodes.Add"),
                    ("G1M4", "omni.graph.nodes.Add"),
                    ("G1Out1", "omni.graph.nodes.Magnitude"),
                    ("G1Out2", "omni.graph.nodes.Magnitude"),
                    ("G2In1", "omni.graph.nodes.ConstantFloat"),
                    ("G2M1", "omni.graph.nodes.Magnitude"),
                    ("G2Out1", "omni.graph.nodes.Magnitude"),
                    ("G3In1", "omni.graph.nodes.ConstantFloat"),
                ],
                keys.CONNECT: [
                    ("G1In1.inputs:value", "G1M1.inputs:a"),
                    ("G1In2.inputs:value", "G1M1.inputs:b"),
                    ("G1M1.outputs:sum", "G1M2.inputs:a"),
                    ("G1M1.outputs:sum", "G1M2.inputs:b"),
                    ("G1M1.outputs:sum", "G1M3.inputs:a"),
                    ("G1M1.outputs:sum", "G1M3.inputs:b"),
                    ("G1M2.outputs:sum", "G1M4.inputs:a"),
                    ("G1M3.outputs:sum", "G1M4.inputs:b"),
                    ("G1M4.outputs:sum", "G1Out1.inputs:input"),
                    ("G1M4.outputs:sum", "G1Out2.inputs:input"),
                    ("G2In1.inputs:value", "G2M1.inputs:input"),
                    ("G2M1.outputs:magnitude", "G2Out1.inputs:input"),
                ],
            },
        )

        prims = tuple(stage.GetPrimAtPath(n.get_prim_path()) for n in nodes)
        g1in1, g1in2, g1m1, _, g1m3, _, _, _, g2in1, g2m1, _, _ = prims

        # Validate the set of visited nodes
        expected = [g1m3, g1m1, g1in1, g1in2, g2m1, g2in1]
        visited = traverse_upstream_graph([g2m1, g1m3])
        actual = [og.Controller.prim(n) for n in visited]
        self.assertCountEqual(expected, actual, "Unexpected set of visited nodes upstream")

    async def test_upstream_connections_can_include_compound_graph_prim(self):
        """Test that calling for uptream connections with compounds can include the compound graph prim."""

        node_type = "omni.graph.nodes.Add"
        controller = og.Controller()
        keys = controller.Keys

        # Test graph set up
        #       |============================|
        # [O1]->|->[I1]---->|I2|----->|I3|-->|--->|O2]
        #       |           |  |             |
        #       |    [I4]------------------->|
        #       |    [I5]->[I6]              |
        #       |============================|
        (_, _, _, nodes) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("Outer1", node_type),
                    ("Outer2", node_type),
                    (
                        "Compound",
                        {
                            keys.CREATE_NODES: [
                                ("Inner1", node_type),
                                ("Inner2", node_type),
                                ("Inner3", node_type),
                                ("Inner4", node_type),
                                ("Inner5", node_type),
                                ("Inner6", node_type),
                            ],
                            keys.PROMOTE_ATTRIBUTES: [
                                ("Inner1.inputs:a", "inputs:value"),
                                ("Inner3.outputs:sum", "outputs:one"),
                                ("Inner4.outputs:sum", "outputs:two"),
                            ],
                            keys.CONNECT: [
                                ("Inner1.outputs:sum", "Inner2.inputs:b"),
                                ("Inner2.outputs:sum", "Inner3.inputs:a"),
                                ("Inner5.outputs:sum", "Inner6.inputs:a"),
                            ],
                        },
                    ),
                ],
                keys.CONNECT: [
                    ("Outer1.outputs:sum", "Compound.inputs:value"),
                    ("Compound.outputs:one", "Outer2.inputs:a"),
                ],
            },
        )

        def assert_valid_traversal(items: List, expected_results: Set[og.Node]):
            """Helper method to validate traversal results."""
            prims = [og.Controller.prim(i) for i in items]

            actual_set = traverse_upstream_graph(prims)
            self.assertEquals(expected_results, actual_set, "Unexpected set of visited nodes downstream")

        def assert_invalid_traversal(items):
            prims = [og.Controller.prim(i) for i in items]
            with self.assertRaises(og.OmniGraphError):
                traverse_upstream_graph(prims)

        compound_graph = nodes["Compound"].get_compound_graph_instance()

        # test traversal with compound graph
        assert_valid_traversal([nodes["Outer2"]], {nodes["Outer1"], nodes["Compound"], nodes["Outer2"]})

        # test with just the compound graph
        assert_valid_traversal([compound_graph], {nodes["Inner1"], nodes["Inner2"], nodes["Inner3"], nodes["Inner4"]})

        # test with a node inside the compound graph
        assert_valid_traversal([compound_graph, nodes["Inner6"]], set(compound_graph.get_nodes()))

        # invalid traversals
        assert_invalid_traversal([nodes["Inner1"], nodes["Outer2"]])  # mixed graphs
        assert_invalid_traversal([compound_graph, nodes["Outer2"]])  # compound and outer node
        assert_invalid_traversal(["/World"])  # invalid node
        assert_invalid_traversal(["/World/Testgraph"])  # non compound graph

    async def test_node_callback_works(self):
        """
        Validate that the node predicate is applied correctly.
        """
        stage = omni.usd.get_context().get_stage()

        # Test graph set up
        #   [G1In1]-->[G1M1]-->[G1M2]--->[G1M4]--->[G1Out1]
        #   [G1In2]---^    |---[G1M3]----|    |--->[G1Out2]
        #
        #   [G2In1]-->[G2M1]-->[G2Out1]
        #
        #   [G3In1]
        controller = og.Controller()
        keys = controller.Keys
        (_, nodes, _, _) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("G1In1", "omni.graph.nodes.ConstantFloat"),
                    ("G1In2", "omni.graph.nodes.ConstantFloat"),
                    ("G1M1", "omni.graph.nodes.Add"),
                    ("G1M2", "omni.graph.nodes.Add"),
                    ("G1M3", "omni.graph.nodes.Add"),
                    ("G1M4", "omni.graph.nodes.Add"),
                    ("G1Out1", "omni.graph.nodes.Magnitude"),
                    ("G1Out2", "omni.graph.nodes.Magnitude"),
                    ("G2In1", "omni.graph.nodes.ConstantFloat"),
                    ("G2M1", "omni.graph.nodes.Magnitude"),
                    ("G2Out1", "omni.graph.nodes.Magnitude"),
                    ("G3In1", "omni.graph.nodes.ConstantFloat"),
                ],
                keys.CONNECT: [
                    ("G1In1.inputs:value", "G1M1.inputs:a"),
                    ("G1In2.inputs:value", "G1M1.inputs:b"),
                    ("G1M1.outputs:sum", "G1M2.inputs:a"),
                    ("G1M1.outputs:sum", "G1M2.inputs:b"),
                    ("G1M1.outputs:sum", "G1M3.inputs:a"),
                    ("G1M1.outputs:sum", "G1M3.inputs:b"),
                    ("G1M2.outputs:sum", "G1M4.inputs:a"),
                    ("G1M3.outputs:sum", "G1M4.inputs:b"),
                    ("G1M4.outputs:sum", "G1Out1.inputs:input"),
                    ("G1M4.outputs:sum", "G1Out2.inputs:input"),
                    ("G2In1.inputs:value", "G2M1.inputs:input"),
                    ("G2M1.outputs:magnitude", "G2Out1.inputs:input"),
                ],
            },
        )

        prims = tuple(stage.GetPrimAtPath(n.get_prim_path()) for n in nodes)
        _, _, _, _, g1m3, _, _, _, _, g2m1, _, _ = prims

        custom_accumulator = []

        def convert_to_prim(node):
            """Convert the visited node to Prim"""
            nonlocal custom_accumulator
            custom_accumulator.append(og.Controller.prim(node))

        # Validate the downstream traversal
        visited = traverse_downstream_graph([g2m1, g1m3], node_callback=convert_to_prim)
        expected = [og.Controller.prim(n) for n in visited]
        self.assertCountEqual(
            expected, custom_accumulator, "Unexpected effect of the node callback in downstream traversal"
        )

        # Validate the upstream traversal
        custom_accumulator = []

        visited = traverse_upstream_graph([g2m1, g1m3], node_callback=convert_to_prim)
        expected = [og.Controller.prim(n) for n in visited]
        self.assertCountEqual(
            expected, custom_accumulator, "Unexpected effect of the node callback in upstream traversal"
        )

    async def test_attribute_predicate_works(self):
        """Validate that the attribute callback allows filtering connections"""
        stage = omni.usd.get_context().get_stage()

        # Test graph set up
        # [OnCustomEvent1]====>[Counter]==================>[Delay]
        # [OnCustomEvent2]=====^       |-->[Negate]        ^
        #                                 [ConstantFloat]--|
        # Where
        #       "==>" : execution connection
        #       "-->" : data connection
        controller = og.Controller()
        keys = controller.Keys
        (_, nodes, _, _) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("OnCustomEvent1", "omni.graph.action.OnCustomEvent"),
                    ("OnCustomEvent2", "omni.graph.action.OnCustomEvent"),
                    ("Counter", "omni.graph.action.Counter"),
                    ("Delay", "omni.graph.action.Delay"),
                    ("Duration", "omni.graph.nodes.ConstantFloat"),
                    ("Negate", "omni.graph.nodes.Negate"),
                ],
                keys.CONNECT: [
                    ("OnCustomEvent1.outputs:execOut", "Counter.inputs:execIn"),
                    ("OnCustomEvent2.outputs:execOut", "Counter.inputs:reset"),
                    ("Counter.outputs:execOut", "Delay.inputs:execIn"),
                    ("Counter.outputs:count", "Negate.inputs:input"),
                    ("Duration.inputs:value", "Delay.inputs:duration"),
                ],
            },
        )

        prims = tuple(stage.GetPrimAtPath(n.get_prim_path()) for n in nodes)
        on_custom_event1, on_custom_event2, counter, delay, _, _ = prims

        def is_execution(attr: og.Attribute) -> bool:
            return attr and attr.get_resolved_type().role == og.AttributeRole.EXECUTION

        visited = traverse_downstream_graph([on_custom_event1], is_execution)
        actual = {og.Controller.prim(n) for n in visited}
        expected = {on_custom_event1, counter, delay}
        self.assertSetEqual(expected, actual, "Unexpected sequence of nodes downstream")

        visited = traverse_upstream_graph([delay], is_execution)
        actual = {og.Controller.prim(n) for n in visited}
        expected = {delay, counter, on_custom_event1, on_custom_event2}
        self.assertSetEqual(expected, actual, "Unexpected sequence of nodes upstream")

    async def test_graph_with_cycles_works(self):
        """Validate that the traversal functions work on graphs with cycles / loops"""
        stage = omni.usd.get_context().get_stage()

        # Test graph set up
        #
        # [OnStageEvent]===>[ForLoop]<======================================
        #                       |  ||============>[Delay]========>[Branch]=||
        #                       |--------------------->[Compare]--^
        #                           [ConstantInt]------^
        #
        # Where
        #       "==>" : execution connection
        #       "-->" : data connection
        controller = og.Controller()
        keys = controller.Keys
        (_, nodes, _, _) = controller.edit(
            "/World/Testgraph",
            {
                keys.CREATE_NODES: [
                    ("OnStageEvent", "omni.graph.action.OnStageEvent"),
                    ("ForLoop", "omni.graph.action.ForLoop"),
                    ("Delay", "omni.graph.action.Delay"),
                    ("ConstantInt", "omni.graph.nodes.ConstantInt"),
                    ("Compare", "omni.graph.nodes.Compare"),
                    ("Branch", "omni.graph.action.Branch"),
                ],
                keys.CONNECT: [
                    ("OnStageEvent.outputs:execOut", "ForLoop.inputs:execIn"),
                    ("ForLoop.outputs:loopBody", "Delay.inputs:execIn"),
                    ("ForLoop.outputs:index", "Compare.inputs:a"),
                    ("Delay.outputs:finished", "Branch.inputs:execIn"),
                    ("ConstantInt.inputs:value", "Compare.inputs:b"),
                    ("Compare.outputs:result", "Branch.inputs:condition"),
                    ("Branch.outputs:execTrue", "ForLoop.inputs:breakLoop"),
                ],
            },
        )

        prims = tuple(stage.GetPrimAtPath(n.get_prim_path()) for n in nodes)
        on_stage_event, for_loop, delay, constant_int, _, branch = prims

        def is_execution(attr: og.Attribute) -> bool:
            return attr and attr.get_resolved_type().role == og.AttributeRole.EXECUTION

        visited = traverse_downstream_graph([delay], is_execution)
        actual = {og.Controller.prim(n) for n in visited}
        expected = {delay, branch, for_loop}
        self.assertSetEqual(expected, actual, "Unexpected sequence of nodes downstream")

        visited = traverse_upstream_graph([delay], is_execution)
        actual = {og.Controller.prim(n) for n in visited}
        expected = {delay, for_loop, branch, on_stage_event}
        self.assertSetEqual(expected, actual, "Unexpected sequence of nodes upstream")

        # The result of a traversal where the starting node has no matching attribute
        # is a set with only the node itself
        visited = traverse_downstream_graph([constant_int], is_execution)
        actual = {og.Controller.prim(n) for n in visited}
        expected = {constant_int}
        self.assertSetEqual(expected, actual, "Unexpected traversal when no connections are valid")
