"""Basic tests of the compute graph"""

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.test
import omni.kit.undo
import omni.usd
import omni.usd.commands


class TestOmniGraphSimple(ogts.OmniGraphTestCase):
    """Run a simple unit test that exercises graph functionality"""

    # ----------------------------------------------------------------------
    async def test_undo_dangling_connections(self):
        # If a dangling connection (see previous test) is created by an undoable
        # command, the connection should be restored if the command is undone.
        (_, (src_node, dest_node), _, _) = og.Controller.edit(
            "/World/PushGraph",
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("src", "omni.graph.tutorials.SimpleData"),
                    ("dest", "omni.graph.tutorials.SimpleData"),
                ],
                og.Controller.Keys.CONNECT: [("src.outputs:a_float", "dest.inputs:a_float")],
            },
        )

        src_attr = src_node.get_attribute("outputs:a_float")
        dest_attr = dest_node.get_attribute("inputs:a_float")

        # Check that the connection is there, on both sides.
        self.assertEqual(len(src_attr.get_downstream_connections()), 1, "Check initial connection, src side.")
        self.assertEqual(len(dest_attr.get_upstream_connections()), 1, "Check initial connection, dest side.")

        # Delete the src prim and confirm that the connection is gone.
        omni.kit.commands.execute("DeletePrims", paths=[src_node.get_prim_path()])
        self.assertEqual(len(dest_attr.get_upstream_connections()), 0, "Check no connection after deletion.")

        # Undo the deletion.
        omni.kit.undo.undo()

        # Wait a couple of ticks for OG to restore everything.
        for _ in range(2):
            await omni.kit.app.get_app().next_update_async()

        # Check that the connection has been restored.
        src_attr = og.Controller.attribute("/World/PushGraph/src.outputs:a_float")
        conns = src_attr.get_downstream_connections()
        self.assertEqual(len(conns), 1, "Check connection exists after undo, src side.")
        self.assertEqual(conns[0], dest_attr, "Check connection is correct after undo, src side.")

        conns = dest_attr.get_upstream_connections()
        self.assertEqual(len(conns), 1, "Check connection exists after undo, dest side.")
        self.assertEqual(conns[0], src_attr, "Check connection is correct after undo, dest side.")

        # Redo the deletion.
        omni.kit.undo.redo()
        self.assertEqual(len(dest_attr.get_upstream_connections()), 0, "Check no connection after redo.")

        # Undo the redo.
        omni.kit.undo.undo()

        # Wait a couple of ticks for OG to restore everything.
        for _ in range(2):
            await omni.kit.app.get_app().next_update_async()

        # Check that the connection has been restored.
        src_attr = og.Controller.attribute("/World/PushGraph/src.outputs:a_float")
        conns = src_attr.get_downstream_connections()
        self.assertEqual(len(conns), 1, "Check connection exists after undo of redo, src side.")
        self.assertEqual(conns[0], dest_attr, "Check connection is correct after undo of redo, src side.")

        conns = dest_attr.get_upstream_connections()
        self.assertEqual(len(conns), 1, "Check connection exists after undo of redo, dest side.")
        self.assertEqual(conns[0], src_attr, "Check connection is correct after undo of redo, src side.")

    # ----------------------------------------------------------------------
    async def test_disconnect(self):
        """
        Check that after a disconnection, we find the default value at input runtime location,
        and that the output still has a correct value
        """
        (_, _, _, _) = og.Controller.edit(
            "/World/PushGraph",
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("src", "omni.graph.tutorials.SimpleData"),
                    ("dest", "omni.graph.tutorials.SimpleData"),
                ],
                og.Controller.Keys.CONNECT: [("src.outputs:a_float", "dest.inputs:a_float")],
            },
        )

        dst_attr = og.Controller.attribute("/World/PushGraph/dest.inputs:a_float")
        src_attr = og.Controller.attribute("/World/PushGraph/src.outputs:a_float")

        await og.Controller.evaluate()

        # since both attribute are connected, they now have the same value: the compute result
        self.assertEqual(dst_attr.get(), 1)
        self.assertEqual(src_attr.get(), 1)

        # do the disconnect
        src_attr.disconnect(dst_attr, True)

        # we get back the default on the input, the output still has the compute result
        self.assertEqual(dst_attr.get(), 0)
        self.assertEqual(src_attr.get(), 1)

    # ----------------------------------------------------------------------
    async def test_unknown_type(self):
        """
        Basic idea of test: Load up an existing file that has an unknown type, and make sure
        that we are not creating a node for it.
        """
        (result, error) = await ogts.load_test_file("TestNodeWithUnknownType.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        await omni.kit.app.get_app().next_update_async()

        graph = og.Controller.graph("/World/Graph")

        unknown_node = graph.get_node("/World/Graph/SomeUnknownNode")
        self.assertFalse(unknown_node.is_valid())

    # ----------------------------------------------------------------------
    async def test_file_format_upgrade(self):
        """
        Basic idea of test: Register a file format upgrade callback.  Load up file with bad node type name, and
        fix it as part of the file format upgrade.
        """

        def my_test_pre_load_file_format_upgrade_callback(_old_format_version, _new_format_version, _graph):
            usd_context = omni.usd.get_context()
            stage = usd_context.get_stage()
            subtract_double_node_prim = stage.GetPrimAtPath("/World/innerPushGraph/test_node_subtract_double_c_")
            self.assertTrue(subtract_double_node_prim.IsValid())
            node_type_attr = subtract_double_node_prim.GetAttribute("node:type")
            node_type_attr.Set("omni.graph.test.SubtractDoubleC")

        handle = og.register_pre_load_file_format_upgrade_callback(my_test_pre_load_file_format_upgrade_callback)
        try:
            (result, error) = await ogts.load_test_file("TestBadNodeTypeName.usda", use_caller_subdirectory=True)
            self.assertTrue(result, error)

            graph = og.get_graph_by_path("/World/innerPushGraph")
            # make sure the box and the attribute we expect are there.
            node = graph.get_node("/World/innerPushGraph/test_node_subtract_double_c_")
            self.assertTrue(node.is_valid())
            self.assertEqual(node.get_node_type().get_node_type(), "omni.graph.test.SubtractDoubleC")
        finally:
            og.deregister_pre_load_file_format_upgrade_callback(handle)

    # ----------------------------------------------------------------------
    async def test_controller_on_initialize(self):
        # Basic idea of test: here we have a test node that uses the og.Controller at initialize
        # Time.  There was a bug that made this not possible.  We test for a clean open here without
        # errors.  See https://nvidia.slack.com/archives/CN05UCXT5/p1614278655049700 for details
        (result, error) = await ogts.load_test_file("TestInitWithController.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

    # ----------------------------------------------------------------------
    async def test_is_compatible(self):
        """
        Test Attribute.is_compatible() for extended and simple types.

        `is_compatible` is verified for a representitive sample of connection combinations,
        bad_connections are connections that should be rejected, ok_connections are connections
        that should be approved.
        Note that this test is meant to take into account (and test!) the auto-conversion mechanism
        """
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("SimpleA", "omni.graph.tutorials.SimpleData"),
                    ("TupleA", "omni.tutorials.TupleData"),
                    ("ArrayA", "omni.graph.tutorials.ArrayData"),
                    ("TupleArrayA", "omni.graph.tutorials.TupleArrays"),
                    ("SimpleB", "omni.graph.tutorials.SimpleData"),
                    ("TupleB", "omni.tutorials.TupleData"),
                    ("ArrayB", "omni.graph.tutorials.ArrayData"),
                    ("TupleArrayB", "omni.graph.tutorials.TupleArrays"),
                    ("ExtendedA", "omni.graph.tutorials.ExtendedTypes"),
                    ("ExtendedB", "omni.graph.tutorials.ExtendedTypes"),
                ]
            },
        )
        await controller.evaluate(graph)

        # Helper for doing verification on a list of connection specs between 2 nodes
        def verify_connections(connections, should_work):
            for src, dest in connections:
                src_node, src_attr = src.split(".", 1)
                dest_node, dest_attr = dest.split(".", 1)
                src_attrib = controller.node(src_node, graph).get_attribute(src_attr)
                dest_attrib = controller.node(dest_node, graph).get_attribute(dest_attr)
                if should_work:
                    self.assertTrue(
                        src_attrib.is_compatible(dest_attrib),
                        f"{src_node}.{src_attr} is compatible with {dest_node}.{dest_attr}",
                    )
                else:
                    self.assertFalse(
                        src_attrib.is_compatible(dest_attrib),
                        f"{src_node}.{src_attr} is not compatible with {dest_node}.{dest_attr}",
                    )

        # Test simple connections
        ok_connections = [
            ("SimpleA.outputs:a_int", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:a_int", "SimpleB.inputs:a_bool"),
            ("SimpleA.outputs:a_int", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:a_int", "SimpleB.inputs:a_int64"),
            ("SimpleA.outputs:unsigned:a_uint", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:unsigned:a_uint", "SimpleB.inputs:a_bool"),
            ("SimpleA.outputs:unsigned:a_uint", "SimpleB.inputs:a_float"),
            ("SimpleA.outputs:unsigned:a_uint", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:unsigned:a_uint", "SimpleB.inputs:a_int64"),
            ("SimpleA.outputs:unsigned:a_uint", "SimpleB.inputs:unsigned:a_uint64"),
            ("SimpleA.outputs:a_float", "SimpleB.inputs:a_float"),
            ("SimpleA.outputs:a_float", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:a_float", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:a_float", "SimpleB.inputs:a_int64"),
            ("SimpleA.outputs:a_double", "SimpleB.inputs:a_float"),
            ("SimpleA.outputs:a_double", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:a_double", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:a_double", "SimpleB.inputs:a_int64"),
            ("SimpleA.outputs:a_int64", "SimpleB.inputs:unsigned:a_uint64"),
            ("SimpleA.outputs:a_int64", "SimpleB.inputs:a_float"),
            ("SimpleA.outputs:a_int64", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:a_int64", "SimpleB.inputs:a_bool"),
            ("SimpleA.outputs:a_int64", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:unsigned:a_uint64"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:a_int64"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:a_float"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:a_bool"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:unsigned:a_uint"),
            ("SimpleA.outputs:unsigned:a_uint64", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:a_half", "SimpleB.inputs:a_half"),
        ]

        bad_connections = [
            ("SimpleA.outputs:a_bool", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:a_bool", "SimpleB.inputs:a_int"),
            ("SimpleA.outputs:a_float", "SimpleB.inputs:unsigned:a_uint"),
            ("SimpleA.outputs:a_int", "SimpleB.inputs:unsigned:a_uint64"),
            ("SimpleA.outputs:a_half", "SimpleB.inputs:a_float"),
            ("SimpleA.outputs:a_double", "SimpleB.inputs:a_half"),
            # FIXME: What about input -> output?
            # ('SimpleB.inputs:a_float', 'SimpleA.outputs:a_float')
        ]
        verify_connections(ok_connections, True)
        verify_connections(bad_connections, False)

        # Test tuple connections
        ok_connections = [
            ("TupleA.outputs:a_float2", "TupleB.inputs:a_double2"),
            ("TupleA.outputs:a_int2", "TupleB.inputs:a_int2"),
            ("TupleA.outputs:a_float2", "TupleB.inputs:a_float2"),
            ("TupleA.outputs:a_double3", "TupleB.inputs:a_double3"),
            ("TupleA.outputs:a_float3", "TupleB.inputs:a_float3"),
        ]

        bad_connections = [
            ("TupleA.outputs:a_double2", "TupleB.inputs:a_double3"),
            ("TupleA.outputs:a_half2", "TupleB.inputs:a_float2"),
            ("TupleA.outputs:a_half2", "SimpleB.inputs:a_half"),
            ("TupleA.outputs:a_double3", "SimpleB.inputs:a_double"),
            ("SimpleA.outputs:a_double", "TupleB.inputs:a_double3"),
        ]
        verify_connections(ok_connections, True)
        verify_connections(bad_connections, False)

        # Test array connections
        ok_connections = [
            ("ArrayA.outputs:result", "ArrayB.inputs:original"),
            ("ArrayA.outputs:negativeValues", "ArrayB.inputs:gates"),
        ]

        bad_connections = [
            ("ArrayA.outputs:result", "ArrayB.inputs:gates"),
            ("ArrayA.outputs:result", "TupleB.inputs:a_float3"),
        ]
        verify_connections(ok_connections, True)
        verify_connections(bad_connections, False)

        # Test tuple-array connections
        ok_connections = [("TupleArrayA.inputs:a", "TupleArrayB.inputs:b")]

        bad_connections = [
            ("TupleArrayA.inputs:a", "ArrayB.inputs:original"),
            ("TupleArrayA.outputs:result", "TupleArrayB.inputs:a"),
        ]
        verify_connections(ok_connections, True)
        verify_connections(bad_connections, False)

        # Test 'any' and enum connections
        ok_connections = [
            ("ExtendedA.outputs:doubledResult", "ExtendedB.inputs:tuple"),
            ("ExtendedA.outputs:flexible", "ExtendedB.inputs:floatOrToken"),
            ("ExtendedA.outputs:tuple", "TupleB.inputs:a_double3"),
            ("ExtendedA.outputs:flexible", "ExtendedB.inputs:flexible"),
        ]
        verify_connections(ok_connections, True)

        # incompatible enum connections
        bad_connections = [
            ("ExtendedA.outputs:flexible", "SimpleB.inputs:a_double"),
            ("SimpleB.outputs:a_double", "ExtendedA.inputs:flexible"),
        ]
        verify_connections(bad_connections, False)

        # Wire up ExtendedA in order to resolve doubledResult to be token, and verify the check fails
        # also resolve outputs:tuple to be int[2], and verify the check fails
        controller.edit(
            "/TestGraph",
            {
                keys.CONNECT: [
                    ("SimpleA.outputs:a_token", "ExtendedA.inputs:floatOrToken"),
                    ("TupleA.outputs:a_int2", "ExtendedA.inputs:tuple"),
                    # FIXME: Do I need to connect the output in order to resolve the type?
                    ("ExtendedA.outputs:doubledResult", "SimpleB.inputs:a_token"),
                ]
            },
        )

        ok_connections = [("TupleA.outputs:a_float2", "ExtendedA.inputs:tuple")]
        verify_connections(ok_connections, True)

        bad_connections = [
            ("ExtendedA.outputs:doubledResult", "SimpleB.inputs:a_float"),
            ("ExtendedA.inputs:floatOrToken", "SimpleB.inputs:a_float"),
            ("SimpleB.inputs:a_float", "ExtendedA.inputs:floatOrToken"),
        ]
        verify_connections(bad_connections, False)

    # ----------------------------------------------------------------------
    async def test_simple_converted_connections(self):
        """
        Test conversions between simple types.

        `simple_converted_connections` tests a representitive sample of connection combinations of different types,
        and makes sure the expected values are actually retrieved
        """
        controller = og.Controller()

        (_, nodes, _, _) = controller.edit(
            "/Root",
            {
                "create_nodes": [
                    ("SimpleA", "omni.graph.tutorials.SimpleData"),
                    ("SimpleB", "omni.graph.tutorials.SimpleData"),
                ],
                "connect": [
                    ("SimpleA.outputs:a_double", "SimpleB.inputs:a_float"),
                    ("SimpleA.outputs:a_int", "SimpleB.inputs:a_bool"),
                    ("SimpleA.outputs:a_int64", "SimpleB.inputs:a_double"),
                    ("SimpleA.outputs:a_float", "SimpleB.inputs:a_int"),
                ],
                "set_values": [
                    ("SimpleA.inputs:a_double", 10.5),
                    ("SimpleA.inputs:a_int", 11),
                    ("SimpleA.inputs:a_int64", 12),
                    ("SimpleA.inputs:a_float", 13.5),
                ],
            },
        )

        await controller.evaluate()
        simple_b_node = nodes[1]
        # A.input.a_double (10.5) -> A.output.a_double (11.5) -> B.input.a_float(11.5) -> B.output.a_float(12.5)
        self.assertEqual(simple_b_node.get_attribute("outputs:a_float").get(), 12.5)
        # A.input.a_int64 (12) -> A.output.a_int64 (13) -> B.input.a_double(13.0) -> B.output.a_double(14.0)
        self.assertEqual(simple_b_node.get_attribute("outputs:a_double").get(), 14.0)
        # A.input.a_float (13.5) -> A.output.a_float (14.5) -> B.input.a_int(14) -> B.output.a_int(15)
        self.assertEqual(simple_b_node.get_attribute("outputs:a_int").get(), 15)
        # A.input.a_int (11) -> A.output.a_int (12) -> B.input.a_bool(True) -> B.output.a_bool(False)
        self.assertFalse(simple_b_node.get_attribute("outputs:a_bool").get())

    # ----------------------------------------------------------------------
    async def test_tuple_converted_connections(self):
        """
        Test conversions between tuple types.

        `tuple_converted_connections` tests a representitive sample of connection combinations of different types,
        and makes sure the expected values are actually retrieved
        """
        controller = og.Controller()

        (_, nodes, _, _) = controller.edit(
            "/Root",
            {
                "create_nodes": [("TupleA", "omni.tutorials.TupleData"), ("TupleB", "omni.tutorials.TupleData")],
                "connect": [
                    ("TupleA.outputs:a_double2", "TupleB.inputs:a_float2"),
                    ("TupleA.outputs:a_int2", "TupleB.inputs:a_double2"),
                    ("TupleA.outputs:a_float2", "TupleB.inputs:a_int2"),
                    ("TupleA.outputs:a_float3", "TupleB.inputs:a_double3"),
                    ("TupleA.outputs:a_double3", "TupleB.inputs:a_float3"),
                ],
                "set_values": [
                    ("TupleA.inputs:a_double2", (10.5, 11.5)),
                    ("TupleA.inputs:a_int2", (13, 14)),
                    ("TupleA.inputs:a_float2", (15.5, 16.5)),
                    ("TupleA.inputs:a_float3", (17.5, 18.5, 19.5)),
                    ("TupleA.inputs:a_double3", (20.5, 21.5, 22.5)),
                ],
            },
        )

        await controller.evaluate()
        tuple_b_node = nodes[1]
        # A.input.a_double2 (10.5,11.5) -> A.output.a_double2 (11.5,12.5)
        #   -> B.input.a_float2(11.5, 12.5) -> B.output.a_float2(12.5,13.5)
        self.assertTrue((tuple_b_node.get_attribute("outputs:a_float2").get() == (12.5, 13.5)).all())
        # A.input.a_int2 (13,14) -> A.output.a_int2 (14,15)
        #   -> B.input.a_double2(14.0, 15.0) -> B.output.a_double2(15.0, 16.0)
        self.assertTrue((tuple_b_node.get_attribute("outputs:a_double2").get() == (15.0, 16.0)).all())
        # A.input.a_float2 (15.5, 16.5) -> A.output.a_float2 (16.5, 17.5)
        #   -> B.input.a_int2(16, 17) -> B.output.a_int2(17, 18)
        self.assertTrue((tuple_b_node.get_attribute("outputs:a_int2").get() == (17, 18)).all())
        # A.input.a_float3 (17.5, 18.5, 19.5) -> A.output.a_float3 (18.5, 19.5, 20.5)
        #   -> B.input.a_double3(18.5, 19.5, 20.5) -> B.output.a_double3(19.5, 20.5, 21.5)
        self.assertTrue((tuple_b_node.get_attribute("outputs:a_double3").get() == (19.5, 20.5, 21.5)).all())
        # A.input.a_double3 (20.5,21.5,22.5) -> A.output.a_double3 (21.5,22.5, 23.5)
        #   -> B.input.a_float3(21.5,22.5, 23.5) -> B.output.a_float3(22.5, 23.5, 24.5)
        self.assertTrue((tuple_b_node.get_attribute("outputs:a_float3").get() == (22.5, 23.5, 24.5)).all())

    # ----------------------------------------------------------------------
    async def test_converted_runtime_connections(self):
        """
        Test conversions of runtime attributes
        ie. when the actual type of the attribute is not listed in the list of accepted input types,
        but a conversion exits
        """

        controller = og.Controller()

        (_, nodes, _, _) = controller.edit(
            "/Root",
            {
                "create_nodes": [
                    ("Simple", "omni.graph.tutorials.SimpleData"),
                    ("Extended", "omni.graph.tutorials.ExtendedTypes"),
                ],
                "connect": [("Simple.outputs:a_double", "Extended.inputs:floatOrToken")],
                "set_values": [
                    ("Simple.inputs:a_double", 1.5),
                    ("Simple.inputs:a_int", 2),
                    ("Simple.inputs:a_float", 3.5),
                ],
            },
        )

        await controller.evaluate()
        ext_node = nodes[1]
        # Simple.in.a_double (1.5) -> Simple.out.a_double (2.5) -> Extended.floatOrToken (2.5) -> Extended.doubled (5.0)
        val = ext_node.get_attribute("outputs:doubledResult").get()
        self.assertTrue(val == 5.0)

    # ----------------------------------------------------------------------
    async def test_converted_prefered_runtime_connections(self):
        """
        Test that the proper conversion of runtime attribute is chosen when several are available
        """

        controller = og.Controller()
        (_, nodes, _, _) = controller.edit(
            "/Root",
            {
                "create_nodes": [
                    ("Simple", "omni.graph.tutorials.SimpleData"),
                    ("A", "omni.graph.nodes.ATan2"),
                    ("A2", "omni.graph.nodes.ATan2"),
                ],
                "connect": [("Simple.outputs:a_int", "A.inputs:a"), ("Simple.outputs:a_int64", "A2.inputs:a")],
            },
        )

        t = nodes[1].get_attribute("inputs:a").get_resolved_type().base_type
        self.assertTrue(t == og.BaseDataType.FLOAT)

        t = nodes[2].get_attribute("inputs:a").get_resolved_type().base_type
        self.assertTrue(t == og.BaseDataType.DOUBLE)

    # ----------------------------------------------------------------------
    async def test_converted_coupled_attributes(self):
        """
        Tests that the conversion for coupled attribute type resolution works
        """

        controller = og.Controller()
        (_, nodes, _, _) = controller.edit(
            "/Root",
            {
                "create_nodes": [
                    ("Mul", "omni.graph.nodes.Multiply"),
                    ("Double", "omni.graph.nodes.ConstantDouble"),
                    ("Float", "omni.graph.nodes.ConstantFloat"),
                ],
                "connect": [("Double.inputs:value", "Mul.inputs:a"), ("Float.inputs:value", "Mul.inputs:b")],
                "set_values": [("Double.inputs:value", 5.0), ("Float.inputs:value", 10.0)],
            },
        )

        await controller.evaluate()

        out = nodes[0].get_attribute("outputs:product")

        # the first pin is used to determine the output type
        base_type = out.get_resolved_type().base_type
        self.assertTrue(base_type == og.BaseDataType.DOUBLE)

        # check that the node executed correctly
        val = out.get()
        self.assertTrue(val == 50)
