# noqa: PLC0302
"""Basic tests of the compute graph"""
import os
import tempfile

import omni.graph.core as og
import omni.graph.core._unstable as ogu
import omni.graph.core.tests as ogts
import omni.graph.tools.ogn as ogn
import omni.kit.commands
import omni.kit.test
import omni.kit.undo
import omni.usd
import omni.usd.commands
from pxr import Gf, OmniGraphSchemaTools, Sdf


# ======================================================================
class SimpleGraph:
    def __init__(self):
        self.graph = None
        self.input_prim = None
        self.simple_node = None
        self.cube_prim = None


# ======================================================================
def create_simple_graph(stage):
    """Create a scene containing an Xform prim, a Cube, and a PushGraph with a SimpleNode."""

    scene = SimpleGraph()

    # create nodes
    path = omni.usd.get_stage_next_free_path(stage, "/" + "InputPrim", True)
    scene.input_prim = stage.DefinePrim(path, "Xform")
    scene.input_prim.CreateAttribute("value", Sdf.ValueTypeNames.Float).Set(0.5)

    (scene.graph, nodes, _, _) = og.Controller.edit(
        "/World/PushGraph",
        {
            og.Controller.Keys.CREATE_NODES: [("SimpleNode", "omni.graph.examples.cpp.Simple")],
            og.Controller.Keys.SET_VALUES: [("SimpleNode.inputs:multiplier", 3.0), ("SimpleNode.inputs:value", 1.0)],
        },
    )
    scene.simple_node = og.Controller.prim(nodes[0])

    path = omni.usd.get_stage_next_free_path(stage, "/" + "Cube", True)
    scene.cube_prim = stage.DefinePrim(path, "Cube")
    scene.cube_prim.CreateAttribute("size", Sdf.ValueTypeNames.Double).Set(2.0)

    return scene


class TestOmniGraphGeneral(ogts.OmniGraphTestCase):

    # ----------------------------------------------------------------------
    async def test_graph_load(self):
        """Test state after a simple graph is loaded from a file"""
        await og.load_example_file("ExampleCube.usda")
        # These nodes are known to be in the file
        self.assertFalse(ogts.verify_node_existence(["/World/graph/NoOp"]))

    # ----------------------------------------------------------------------
    async def test_graph_creation(self):
        """Test state after manually creating a simple graph"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        scene = create_simple_graph(stage)
        self.assertIsNotNone(scene, "Failed to create a simple graph")
        self.assertIsNotNone(scene.input_prim, "Failed to create input primitive in simple graph")
        self.assertIsNotNone(scene.simple_node, "Failed to create simple node in simple graph")
        self.assertIsNotNone(scene.cube_prim, "Failed to create cube primitive in simple graph")
        await omni.kit.app.get_app().next_update_async()
        ogts.dump_graph()

        # Verify that the simple_node prim belongs to an actual OG node in the graph.
        node_prims = [og.Controller.prim(node) for node in scene.graph.get_nodes()]
        self.assertTrue(scene.simple_node in node_prims)

    # ----------------------------------------------------------------------
    async def test_graph_disable(self):
        """Exercise disabling a graph"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        scene = create_simple_graph(stage)
        controller = og.Controller()
        graph = scene.graph
        simple_node = og.get_node_by_path(scene.simple_node.GetPath().pathString)
        multiplier_attr = simple_node.get_attribute("inputs:multiplier")
        in_value_attr = simple_node.get_attribute("inputs:value")
        out_value_attr = simple_node.get_attribute("outputs:value")

        # set some values on simple node
        multiplier_attr.set(2)
        in_value_attr.set(3)

        # test the graph works
        await controller.evaluate()
        self.assertEqual(out_value_attr.get(), 6)

        # test the graph works with an update
        multiplier_attr.set(3)
        await controller.evaluate()
        self.assertEqual(out_value_attr.get(), 9)

        # disable the graph
        graph.set_disabled(True)

        # test the graph stopped taking into account updates
        multiplier_attr.set(4)
        await controller.evaluate()
        self.assertEqual(out_value_attr.get(), 9)

        # re-enable the graph
        graph.set_disabled(False)

        # value should now update
        await controller.evaluate()
        self.assertEqual(out_value_attr.get(), 12)

    # ----------------------------------------------------------------------
    async def test_usd_update_settings(self):
        """
        Basic idea of test: here we test the USD write back policy of OmniGraph under normal circumstances.
        We get the USD prim that is backing an intermediate node, and test all possible settings against it to make sure
        it is updating only when expected.
        """
        (result, error) = await ogts.load_test_file("TestTranslatingCube.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        subtract_node_prim = stage.GetPrimAtPath("/World/moveX/subtract")
        output_usd_attr = subtract_node_prim.GetAttribute("outputs:out")

        # test that there is no change in the intermediate USD value:
        # ------------------------------------------

        t0 = output_usd_attr.Get()
        await omni.kit.app.get_app().next_update_async()
        t1 = output_usd_attr.Get()

        diff = abs(t1 - t0)
        self.assertAlmostEqual(diff, 0.0, places=3)

    # ----------------------------------------------------------------------
    async def test_reload_from_stage(self):
        # The basic idea of the test is as follows:
        # We load up the TestTranslatingCube.usda file, which has a cube
        # that continually moves in X.  We make sure first the
        # node is working, and then we reload the graph from the stage.  Once we do
        # that, we make sure that the graph is still working - make sure the cube still moves
        (result, error) = await ogts.load_test_file("TestTranslatingCube.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        # make sure the box and the attribute we expect are there.
        graph = og.get_graph_by_path("/World/moveX")
        cube_prim = stage.GetPrimAtPath("/World/Cube")
        translate_attr = cube_prim.GetAttribute("xformOp:translate")
        self.assertTrue(translate_attr.IsValid())

        # get the initial position of the cube
        translate = translate_attr.Get()
        t0 = translate[0]

        # wait 1 frame
        await omni.kit.app.get_app().next_update_async()

        # get the updated cube position and make sure it moved
        translate = translate_attr.Get()
        t1 = translate[0]

        diff = abs(t1 - t0)
        self.assertGreater(diff, 0.1)

        # subtract_node = graph.get_node("/World/moveX/subtract")
        # node_handle = subtract_node.get_handle()

        # now reload the graph and make sure the node has a new handle (showing that it's been deleted and recreated)
        # and that everything still works.
        graph.reload_from_stage()

        # Note that we have to re-fetch a lot of the constructs, as those would have been destroyed and recreated during
        # the graph reload process:
        # subtract_node = graph.get_node("/World/moveX/subtract")
        # new_node_handle = subtract_node.get_handle()
        # Almost certainly the new handle will not be the same as the old one, but we're not doing the check here
        # because it possibly could be the same, and that could lead to flaky test.
        # self.assertTrue(new_node_handle != node_handle)

        cube_prim = stage.GetPrimAtPath("/World/Cube")
        translate_attr = cube_prim.GetAttribute("xformOp:translate")
        self.assertTrue(translate_attr.IsValid())
        # get the initial position of the cube
        translate = translate_attr.Get()
        t0 = translate[0]

        # wait 1 frame
        await omni.kit.app.get_app().next_update_async()

        # get the updated cube position and make sure it moved
        translate = translate_attr.Get()
        t1 = translate[0]

        diff = abs(t1 - t0)
        self.assertGreater(diff, 0.1)

    # ----------------------------------------------------------------------
    async def test_reload_from_stage_retains_instances(self):
        """Validates that the reload from stage call will retain any instances that may be referencing the node"""

        graph_path = "/World/Graph"
        num_instances = 3
        controller = og.Controller(update_usd=True)
        keys = controller.Keys

        # Simple graph that increments a counter attribute on the graph target
        # [ReadPrim]->[Add 1]->[WritePrim]
        (graph, _, _, nodes) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadPrim", "omni.graph.nodes.ReadPrimAttribute"),
                    ("WritePrim", "omni.graph.nodes.WritePrimAttribute"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("Constant", "omni.graph.nodes.ConstantDouble"),
                    ("GraphTarget", "omni.graph.nodes.GetGraphTargetPrim"),
                ],
                keys.SET_VALUES: [
                    ("Constant.inputs:value", 1.0),
                    ("ReadPrim.inputs:name", "counter"),
                    ("WritePrim.inputs:name", "counter"),
                ],
                keys.CONNECT: [
                    ("GraphTarget.outputs:prim", "ReadPrim.inputs:prim"),
                    ("GraphTarget.outputs:prim", "WritePrim.inputs:prim"),
                    ("ReadPrim.outputs:value", "Add.inputs:a"),
                    ("Constant.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "WritePrim.inputs:value"),
                ],
            },
        )

        # create instances
        for i in range(0, num_instances):
            prim_name = f"/World/Prim_{i}"
            prim = omni.usd.get_context().get_stage().DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(omni.usd.get_context().get_stage(), prim_name, graph_path)
            prim.CreateAttribute("counter", Sdf.ValueTypeNames.Int).Set(0)

        await controller.evaluate()
        await omni.kit.app.get_app().next_update_async()

        # capture the existing values
        attr = nodes["Add"].get_attribute("outputs:sum")
        values = []
        for i in range(0, num_instances):
            values.append(attr.get(instance=i))
        attr_path = attr.get_path()

        # reload the graph here. Note this will invalidate any handles to og objects (except the graph)
        graph.reload_from_stage()

        # reevaluate
        await controller.evaluate()
        await omni.kit.app.get_app().next_update_async()
        attr = og.Controller.attribute(attr_path)

        # validate the values have been updated
        for i in range(0, num_instances):
            self.assertGreater(attr.get(instance=i), values[i])

    # ----------------------------------------------------------------------
    async def test_invalid_graph_objects(self):
        # The basic idea of the test is as follows:
        # Often the graph constructs like node, attribute, graph etc. will outlive their python wrapers.
        # Accessing these deleted things will lead to a crash.  We have put in measures to prevent things
        # from crashing, even if the underlying objects have been destroyed.  This test tries out these
        # scenarios.  This test also tests the deletion of global graphs, by deleting the corresponding
        # wrapper node in the orchestration graph
        (result, error) = await ogts.load_test_file("TestTranslatingCube.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph = og.Controller.graph("/World/moveX")

        # make sure the box and the attribute we expect are there.
        context = graph.get_default_graph_context()
        subtract_node = graph.get_node("/World/moveX/subtract")
        output_attr = subtract_node.get_attribute("outputs:out")
        self.assertTrue(output_attr.is_valid())
        self.assertTrue(subtract_node.is_valid())
        graph.destroy_node("/World/moveX/subtract", True)
        self.assertFalse(subtract_node.is_valid())
        self.assertFalse(output_attr.is_valid())
        self.assertTrue(context.is_valid())
        self.assertTrue(graph.is_valid())

        orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION
        )
        self.assertEqual(len(orchestration_graphs), 1)
        orchestration_graph = orchestration_graphs[0]
        global_graph_nodes = orchestration_graph.get_nodes()
        self.assertTrue(len(global_graph_nodes) == 1)
        global_implicit_graph_node = global_graph_nodes[0]
        orchestration_graph.destroy_node(global_graph_nodes[0].get_prim_path(), False)
        self.assertFalse(global_implicit_graph_node.is_valid())
        self.assertFalse(context.is_valid())
        self.assertFalse(graph.is_valid())

        # Do not use something like new stage to destroy the objects and assert invalidity.  This can cause
        # a race condition where the destruction of the existing stuff on the stage is not complete when
        # the new stage is ready
        # await omni.usd.get_context().new_stage_async()

        # self.assertFalse(context.is_valid())
        # self.assertFalse(graph.is_valid())
        # self.assertFalse(output_attr.is_valid())

    # ----------------------------------------------------------------------
    async def test_memory_type_and_optional_keywords(self):
        """Test the persistence of the memoryType and optional keyword values in .ogn nodes"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # Create some nodes with known memoryType and optional values
        #
        #  C++ Nodes
        #    omni.graph.tutorials.CpuGpuData
        #        node memory type = any
        #        inputs:is_gpu memory type = cpu
        #        inputs:a memory type = any (inherited from node)
        #    omni.graph.tutorials.CudaData
        #        node memory type = cuda
        #        inputs:a memory type = cuda (inherited from node)
        #        inputs:multiplier memory type = cpu
        #    omni.graph.nodes.CopyAttr
        #        inputs:fullData optional = false
        #        inputs:partialData optional = true
        #
        #  Python Nodes
        #    omni.graph.tutorials.CpuGpuExtendedPy
        #        node memory type = cpu
        #        inputs:cpuData memory type = any
        #        inputs:gpu memory type = cpu (inherited from node)
        #    omni.graph.test.PerturbPointsGpu
        #        node memory type = cuda
        #        inputs:percentModified memory type = cpu
        #        inputs:minimum memory type = cuda (inherited from node)
        #    omni.graph.tutorials.TutorialSimpleDataPy
        #        inputs:a_bool optional = true
        #        inputs:a_half optional = false
        (
            _,
            (
                cpu_gpu_node,
                cuda_node,
                optional_node,
                cpu_gpu_py_node,
                cuda_py_node,
                optional_py_node,
            ),
            _,
            _,
        ) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("CpuGpu", "omni.graph.tutorials.CpuGpuData"),
                    ("Cuda", "omni.graph.tutorials.CudaData"),
                    ("Optional", "omni.graph.nodes.CopyAttribute"),
                    ("CpuGpuPy", "omni.graph.tutorials.CpuGpuBundlesPy"),
                    ("CudaPy", "omni.graph.test.PerturbPointsGpu"),
                    ("OptionalPy", "omni.graph.tutorials.SimpleDataPy"),
                ]
            },
        )
        memory_type_key = ogn.MetadataKeys.MEMORY_TYPE

        self.assertEqual(cpu_gpu_node.get_node_type().get_metadata(memory_type_key), ogn.MemoryTypeValues.ANY)
        is_gpu_attr = cpu_gpu_node.get_attribute("inputs:is_gpu")
        self.assertEqual(is_gpu_attr.get_metadata(memory_type_key), ogn.MemoryTypeValues.CPU)
        a_attr = cpu_gpu_node.get_attribute("inputs:a")
        self.assertEqual(a_attr.get_metadata(memory_type_key), None)

        self.assertEqual(cuda_node.get_node_type().get_metadata(memory_type_key), ogn.MemoryTypeValues.CUDA)
        multiplier_attr = cuda_node.get_attribute("inputs:multiplier")
        self.assertEqual(multiplier_attr.get_metadata(memory_type_key), ogn.MemoryTypeValues.CPU)
        a_attr = cuda_node.get_attribute("inputs:a")
        self.assertEqual(a_attr.get_metadata(memory_type_key), None)

        self.assertEqual(optional_node.get_node_type().get_metadata(memory_type_key), None)
        full_data_attr = optional_node.get_attribute("inputs:fullData")
        self.assertFalse(full_data_attr.is_optional_for_compute)
        partial_data_attr = optional_node.get_attribute("inputs:partialData")
        self.assertTrue(partial_data_attr.is_optional_for_compute)

        self.assertEqual(cpu_gpu_py_node.get_node_type().get_metadata(memory_type_key), None)
        gpu_attr = cpu_gpu_py_node.get_attribute("inputs:cpuBundle")
        self.assertEqual(gpu_attr.get_metadata(memory_type_key), None)
        cpu_data_attr = cpu_gpu_py_node.get_attribute("inputs:gpuBundle")
        self.assertEqual(cpu_data_attr.get_metadata(memory_type_key), ogn.MemoryTypeValues.CUDA)
        cpu_data_attr = cpu_gpu_py_node.get_attribute("outputs_cpuGpuBundle")
        self.assertEqual(cpu_data_attr.get_metadata(memory_type_key), ogn.MemoryTypeValues.ANY)

        self.assertEqual(cuda_py_node.get_node_type().get_metadata(memory_type_key), ogn.MemoryTypeValues.CUDA)
        percent_attr = cuda_py_node.get_attribute("inputs:percentModified")
        self.assertEqual(percent_attr.get_metadata(memory_type_key), ogn.MemoryTypeValues.CPU)
        minimum_attr = cuda_py_node.get_attribute("inputs:minimum")
        self.assertEqual(minimum_attr.get_metadata(memory_type_key), None)

        self.assertEqual(optional_py_node.get_node_type().get_metadata(memory_type_key), None)
        a_half_attr = optional_py_node.get_attribute("inputs:a_half")
        self.assertEqual(a_half_attr.is_optional_for_compute, False)
        a_bool_attr = optional_py_node.get_attribute("inputs:a_bool")
        self.assertEqual(a_bool_attr.is_optional_for_compute, True)

        # Check that the three methods of getting Cpu pointers to Gpu data give the same answer
        # Method 1: Directly from the attribute value helper
        data_view = og.AttributeValueHelper(cuda_node.get_attribute("outputs:points"))
        data_view.gpu_ptr_kind = og.PtrToPtrKind.CPU
        data_view_data = data_view.get(on_gpu=True)

        # Method 2: From the controller class get() method
        controller_data = og.Controller.get(
            cuda_node.get_attribute("outputs:points"), on_gpu=True, gpu_ptr_kind=og.PtrToPtrKind.CPU
        )
        self.assertEqual(data_view_data.memory, controller_data.memory)

        # Method 3: From the controller object get() method
        controller = og.Controller(cuda_node.get_attribute("outputs:points"))
        controller.gpu_ptr_kind = og.PtrToPtrKind.CPU
        controller_data = controller.get(on_gpu=True)
        self.assertEqual(data_view_data.memory, controller_data.memory)

    # ----------------------------------------------------------------------
    async def test_equality_comparisons(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        scene = create_simple_graph(stage)
        graph = scene.graph

        # Nodes
        simple_node = og.get_node_by_path(scene.simple_node.GetPath().pathString)
        self.assertEqual(simple_node, simple_node)
        self.assertNotEqual(simple_node, None)

        simple_node2 = graph.get_node(simple_node.get_prim_path())
        self.assertEqual(simple_node, simple_node2)

        cube_node = og.get_node_by_path(scene.cube_prim.GetPath().pathString)
        self.assertNotEqual(simple_node, cube_node)

        # TODO: These emit intentional error messages which the build
        #       framework picks up as test failures. Need to find another
        #       way to test this.
        #
        # bad_node = graph.get_node('bad_node')
        # bad_node2 = graph.get_node('bad_node2')
        # self.assertNotEqual(simple_node, bad_node)
        # self.assertEqual(bad_node, bad_node2)

        # Graphs
        self.assertEqual(graph, graph)
        self.assertNotEqual(graph, None)

        graph2 = simple_node.get_graph()
        self.assertEqual(graph, graph2)

        # TODO: These emit intentional error messages which the build
        #       framework picks up as test failures. Need to find another
        #       way to test this.
        #
        # bad_graph = bad_node.get_graph()
        # bad_graph2 = bad_node2.get_graph()
        # self.assertNotEqual(graph, bad_graph)
        # self.assertEqual(bad_graph, bad_graph2)

        # Attributes
        multiplier_attr = simple_node.get_attribute("inputs:multiplier")
        self.assertEqual(multiplier_attr, multiplier_attr)
        self.assertNotEqual(multiplier_attr, None)

        multiplier_attr2 = simple_node.get_attribute("inputs:multiplier")
        self.assertEqual(multiplier_attr, multiplier_attr2)

        value_attr = simple_node.get_attribute("inputs:value")
        self.assertNotEqual(multiplier_attr, value_attr)

        # TODO: These emit intentional error messages which the build
        #       framework picks up as test failures. Need to find another
        #       way to test this.
        #
        # bad_attr = bad_node.get_attribute('bad_attr')
        # bad_attr2 = bad_node.get_attribute('bad_attr2')
        # self.assertNotEqual(multiplier_attr, bad_attr)
        # self.assertEqual(bad_attr, bad_attr2)

        # AttributeData
        mult_data = multiplier_attr.get_attribute_data()
        self.assertEqual(mult_data, mult_data)
        self.assertNotEqual(mult_data, None)

        mult_data2 = multiplier_attr.get_attribute_data()
        self.assertEqual(mult_data, mult_data2)

        value_data = value_attr.get_attribute_data()
        self.assertNotEqual(mult_data, value_data)

        # Types
        t1 = og.Type(og.INT, 2, 1, og.COLOR)
        same_as_t1 = og.Type(og.INT, 2, 1, og.COLOR)
        t2 = og.Type(og.FLOAT, 2, 1, og.COLOR)
        t3 = og.Type(og.INT, 3, 1, og.COLOR)
        t4 = og.Type(og.INT, 2, 0, og.COLOR)
        t5 = og.Type(og.INT, 2, 1, og.POSITION)
        self.assertEqual(t1, t1)
        self.assertNotEqual(t1, None)
        self.assertEqual(t1, same_as_t1)
        self.assertNotEqual(t1, t2)
        self.assertNotEqual(t1, t3)
        self.assertNotEqual(t1, t4)
        self.assertNotEqual(t1, t5)

    # ----------------------------------------------------------------------
    async def test_hashability(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        scene = create_simple_graph(stage)
        graph = scene.graph

        # Nodes
        simple_node = og.get_node_by_path(scene.simple_node.GetPath().pathString)
        simple_node2 = graph.get_node(simple_node.get_prim_path())
        cube_node = og.get_node_by_path(scene.cube_prim.GetPath().pathString)
        d = {simple_node: 1, cube_node: 2, simple_node2: 3}
        self.assertEqual(len(d), 2)
        self.assertEqual(d[simple_node], 3)

        # Graphs
        graph2 = simple_node.get_graph()

        d = {graph: -3, graph2: -17}
        self.assertEqual(len(d), 1)
        self.assertEqual(d[graph], -17)

        # Attributes
        multiplier_attr = simple_node.get_attribute("inputs:multiplier")
        multiplier_attr2 = simple_node.get_attribute("inputs:multiplier")
        value_attr = simple_node.get_attribute("inputs:value")

        d = {multiplier_attr: 5, value_attr: 6, multiplier_attr2: 7}
        self.assertEqual(len(d), 2)
        self.assertEqual(d[multiplier_attr], 7)

        # AttributeData
        mult_data = multiplier_attr.get_attribute_data()
        mult_data2 = multiplier_attr.get_attribute_data()
        value_data = value_attr.get_attribute_data()

        d = {mult_data: 8, value_data: 9, mult_data2: 10}
        self.assertEqual(len(d), 2)
        self.assertEqual(d[mult_data], 10)

        # Types
        t1 = og.Type(og.INT, 2, 1, og.COLOR)
        same_as_t1 = og.Type(og.INT, 2, 1, og.COLOR)
        t2 = og.Type(og.FLOAT, 2, 1, og.COLOR)
        t3 = og.Type(og.INT, 3, 1, og.COLOR)
        t4 = og.Type(og.INT, 2, 0, og.COLOR)
        t5 = og.Type(og.INT, 2, 1, og.POSITION)

        d = {t1: 1, t2: 2, t3: 3, t4: 5, t5: 6, same_as_t1: 4}
        self.assertEqual(len(d), 5)
        self.assertEqual(d[t1], 4)

    # ----------------------------------------------------------------------
    async def test_dirty_push_time_change(self):
        """Test that dirty_push evaluator dirties the time node when time changes"""
        controller = og.Controller()
        (_, (read_time, sub), _, _) = controller.edit(
            {"graph_path": "/World/TestGraph", "evaluator_name": "dirty_push"},
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("ReadTime", "omni.graph.nodes.ReadTime"),
                    ("Sub", "omni.graph.test.SubtractDoubleC"),
                ],
                og.Controller.Keys.CONNECT: [
                    ("ReadTime.outputs:frame", "Sub.inputs:a"),
                    ("ReadTime.outputs:frame", "Sub.inputs:b"),
                ],
            },
        )
        await controller.evaluate()

        # Verify when anim time does not change, the downstream node does not compute
        cc_0 = sub.get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        cc_1 = sub.get_compute_count()

        # FIXME: OM-45551: Assert fails due to a bug in the lazy evaluator
        # self.assertEqual(cc_0, cc_1)

        # Verify when anim time DOES change, the downstream node DOES compute
        time_out = controller.attribute("outputs:time", read_time)

        # Start playback
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_start_time(1.0)
        timeline.set_end_time(10.0)
        timeline.set_target_framerate(timeline.get_time_codes_per_seconds())
        timeline.play()

        # One update and record start values
        await omni.kit.app.get_app().next_update_async()
        cc_0 = sub.get_compute_count()
        t_0 = og.Controller.get(time_out)
        self.assertAlmostEqual(t_0, 1.0, places=1)

        # One update and compare values
        await omni.kit.app.get_app().next_update_async()

        cc_1 = sub.get_compute_count()
        t_1 = og.Controller.get(time_out)

        self.assertGreater(t_1, t_0)
        self.assertEqual(cc_1, cc_0 + 1)

    # ----------------------------------------------------------------------
    async def test_attribute_deprecation(self):
        """Test the attribute deprecation API."""
        controller = og.Controller()
        (_, nodes, _, _) = controller.edit(
            {"graph_path": "/World/TestGraph", "evaluator_name": "dirty_push"},
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("cpp_node1", "omni.graph.test.ComputeErrorCpp"),
                    ("cpp_node2", "omni.graph.test.ComputeErrorCpp"),
                    ("py_node1", "omni.graph.test.ComputeErrorPy"),
                    ("py_node2", "omni.graph.test.ComputeErrorPy"),
                ]
            },
        )
        for node in nodes:
            deprecated_in_ogn = node.get_attribute("inputs:deprecatedInOgn")
            self.assertTrue(deprecated_in_ogn.is_deprecated())
            self.assertEqual(deprecated_in_ogn.deprecation_message(), "Use 'dummyIn' instead.")

            deprecated_in_init = node.get_attribute("inputs:deprecatedInInit")
            self.assertTrue(deprecated_in_init.is_deprecated())
            self.assertEqual(deprecated_in_init.deprecation_message(), "Use 'dummyIn' instead.")

            dummy_in = node.get_attribute("inputs:dummyIn")
            self.assertFalse(dummy_in.is_deprecated())
            self.assertEqual(dummy_in.deprecation_message(), "")

            # Deprecate the dummyIn attr.
            og._internal.deprecate_attribute(dummy_in, "Deprecated on the fly.")  # noqa: PLW0212
            self.assertTrue(dummy_in.is_deprecated())
            self.assertEqual(dummy_in.deprecation_message(), "Deprecated on the fly.")

    # ----------------------------------------------------------------------
    async def test_minimal_population(self):
        """Test that when fabric minimal population is enabled, that we still initialize graphs properly.
        See OM-108844 for more details"""
        controller = og.Controller()
        (_, nodes, _, _) = controller.edit(
            {"graph_path": "/World/TestGraph"},
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("find_prims", "omni.graph.nodes.FindPrims"),
                ]
            },
        )
        await controller.evaluate()

        self.assertEqual(nodes[0].get_attribute("inputs:type").get(), "*")

        with tempfile.TemporaryDirectory() as test_directory:
            # Need to set the default prim, otherwise creating a reference/payload to this file will cause an error
            stage = omni.usd.get_context().get_stage()
            stage.SetDefaultPrim(stage.GetPrimAtPath("/World"))

            tmp_file_path = os.path.join(test_directory, "tmp_test_minimal_population.usda")
            await omni.usd.get_context().save_as_stage_async(tmp_file_path)

            # Run a usdrt query to ensure that fabric minimal population is enabled
            from usdrt import Usd as UsdRtUsd

            usdrt_stage = UsdRtUsd.Stage.Attach(omni.usd.get_context().get_stage_id())
            self.assertEqual(usdrt_stage.GetPrimsWithTypeName("OmniGraph"), ["/World/TestGraph"])

            # Test adding as payload
            await omni.usd.get_context().new_stage_async()
            await omni.kit.app.get_app().next_update_async()

            # Run a usdrt query to ensure that fabric minimal population is enabled
            from usdrt import Usd as UsdRtUsd

            usdrt_stage = UsdRtUsd.Stage.Attach(omni.usd.get_context().get_stage_id())
            self.assertEqual(usdrt_stage.GetPrimsWithTypeName("OmniGraph"), [])

            stage = omni.usd.get_context().get_stage()
            payload = stage.DefinePrim("/payload")
            payload.GetPayloads().AddPayload(tmp_file_path)
            await omni.kit.app.get_app().next_update_async()

            node = controller.node("/payload/TestGraph/find_prims")
            self.assertEqual(node.get_attribute("inputs:type").get(), "*")

            # Test adding as reference
            await omni.usd.get_context().new_stage_async()
            await omni.kit.app.get_app().next_update_async()
            stage = omni.usd.get_context().get_stage()

            ref = stage.DefinePrim("/ref")
            ref.GetReferences().AddReference(tmp_file_path)
            await omni.kit.app.get_app().next_update_async()

            node = controller.node("/ref/TestGraph/find_prims")
            self.assertEqual(node.get_attribute("inputs:type").get(), "*")


# ===================================================================
class TestOmniGraphSimpleMisc(ogts.OmniGraphTestCase):
    """Misc tests that issue expected errors functionality"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__cb_count = 0
        self.__error_node = None

    async def test_compute_messages(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        scene = create_simple_graph(stage)
        graph = scene.graph
        graph.set_disabled(True)  # to ensure "compute" messages are managed only by hand in the first tests
        node = og.get_node_by_path(scene.simple_node.GetPath().pathString)

        # A new node should start with no messages.
        self.assertEqual(len(node.get_compute_messages(og.INFO)), 0)
        self.assertEqual(len(node.get_compute_messages(og.WARNING)), 0)
        self.assertEqual(len(node.get_compute_messages(og.ERROR)), 0)

        # Check counts.
        msg1 = "Ignore this error/warning message #1"
        msg2 = "Ignore this error/warning message #2"
        msg3 = "Ignore this error/warning message #3"
        node.log_compute_message(og.INFO, msg1)
        node.log_compute_message(og.WARNING, msg1)
        node.log_compute_message(og.WARNING, msg2)
        self.assertEqual(len(node.get_compute_messages(og.INFO)), 1)
        self.assertEqual(len(node.get_compute_messages(og.WARNING)), 2)

        # Check that logging detects duplicates.
        self.assertFalse(node.log_compute_message(og.ERROR, msg2))
        self.assertFalse(node.log_compute_message(og.ERROR, msg1))
        self.assertTrue(node.log_compute_message(og.ERROR, msg2))

        # Check that duplicates are only recorded once per severity level.
        msgs = node.get_compute_messages(og.ERROR)
        self.assertEqual(len(msgs), 2)

        # Check message retrieval. Note that while the current implementation
        # does preserve the order in which the messages were issued, that's
        # not a requirement, so we only test that we get each message back once.
        counts = {msg1: 0, msg2: 0}
        for msg in msgs:
            # The logging ABI adds context information to the message so we can't
            # do an exact comparison with our original, only that the logged message
            # contains our original.
            if msg.find(msg1) != -1:
                counts[msg1] += 1
            elif msg.find(msg2) != -1:
                counts[msg2] += 1

        self.assertEqual(list(counts.values()).count(1), 2)

        # Repeat a message at a new compute count and make sure the message
        # count hasn't changed.
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(node.log_compute_message(og.ERROR, msg1))
        self.assertEqual(len(node.get_compute_messages(og.ERROR)), 2)

        # Add a new (non-repeated) message at the new compute count.
        self.assertFalse(node.log_compute_message(og.ERROR, msg3))

        # If we clear old messages all but the one we just repeated and
        # the one just added should be removed.
        self.assertEqual(node.clear_old_compute_messages(), 4)
        self.assertEqual(len(node.get_compute_messages(og.INFO)), 0)
        self.assertEqual(len(node.get_compute_messages(og.WARNING)), 0)
        self.assertEqual(len(node.get_compute_messages(og.ERROR)), 2)

        # Clear everything.
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(node.clear_old_compute_messages(), 2)
        self.assertEqual(len(node.get_compute_messages(og.INFO)), 0)
        self.assertEqual(len(node.get_compute_messages(og.WARNING)), 0)
        self.assertEqual(len(node.get_compute_messages(og.ERROR)), 0)

        # Ensure it still all works after clearing.
        self.assertFalse(node.log_compute_message(og.ERROR, msg1))
        self.assertEqual(len(node.get_compute_messages(og.ERROR)), 1)
        self.assertTrue(node.get_compute_messages(og.ERROR)[0].find(msg1) != -1)

        # Test with nodes which will generate messages during compute.
        graph.set_disabled(False)
        for node_name, node_type in [
            ("CppErrorNode", "omni.graph.test.ComputeErrorCpp"),
            ("PyErrorNode", "omni.graph.test.ComputeErrorPy"),
        ]:

            (_, nodes, _, _) = og.Controller.edit(graph, {og.Controller.Keys.CREATE_NODES: [(node_name, node_type)]})
            self.assertEqual(len(nodes), 1)
            self.__error_node = nodes[0]

            # It should start with no messages.
            self.assertEqual(len(self.__error_node.get_compute_messages(og.INFO)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.WARNING)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.ERROR)), 0)

            # Set it to generate a warning.
            severity_attr = self.__error_node.get_attribute("inputs:severity")
            message_attr = self.__error_node.get_attribute("inputs:message")
            msg4 = "Ignore this error/warning message #4"

            message_attr.set(msg4)
            severity_attr.set("warning")

            # There should be no warning yet since the node has not evaluated.
            self.assertEqual(len(self.__error_node.get_compute_messages(og.INFO)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.WARNING)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.ERROR)), 0)

            # Evaluate the graph and check again.
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(len(self.__error_node.get_compute_messages(og.INFO)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.ERROR)), 0)
            msgs = self.__error_node.get_compute_messages(og.WARNING)
            self.assertEqual(len(msgs), 1)
            self.assertTrue(msgs[0].find(msg4) != -1)

            # Have it issue a different error message. After evaluation we should only
            # have the error not the warning.
            msg5 = "Ignore this error/warning message #5"
            message_attr.set(msg5)
            severity_attr.set("error")
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(len(self.__error_node.get_compute_messages(og.INFO)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.WARNING)), 0)
            msgs = self.__error_node.get_compute_messages(og.ERROR)
            self.assertEqual(len(msgs), 1)
            self.assertTrue(msgs[0].find(msg5) != -1)

            # Stop the node generating errors. After the next evaluate there
            # should be no messages.
            severity_attr.set("none")
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(len(self.__error_node.get_compute_messages(og.INFO)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.WARNING)), 0)
            self.assertEqual(len(self.__error_node.get_compute_messages(og.ERROR)), 0)

            # Set up a callback for changes in the node's error status.
            self.__cb_count = 0

            def error_status_changed(nodes, graph):
                self.assertEqual(len(nodes), 1)
                self.assertEqual(nodes[0], self.__error_node)
                self.__cb_count += 1

            cb = graph.register_error_status_change_callback(error_status_changed)

            # There are currently no errors so the callback should not get
            # called.
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(self.__cb_count, 0)

            # Setting a warning message should call the callback.
            message_attr.set(msg4)
            severity_attr.set("warning")
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(self.__cb_count, 1)

            # Converting the warning to an error should call the callback.
            severity_attr.set("error")
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(self.__cb_count, 2)

            # Setting the same message again should not call the callback.
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(self.__cb_count, 2)

            # A clean compute should clear the message and call the callback.
            severity_attr.set("none")
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(self.__cb_count, 3)

            # No more messages so further evaluations should not call the callback.
            await omni.kit.app.get_app().next_update_async()
            self.assertEqual(self.__cb_count, 3)

            graph.deregister_error_status_change_callback(cb)

    async def test_write_to_backing(self):
        """Test the GraphContext.write_to_backing functionality"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        (graph, (constant_string,), _, path_to_object_map) = og.Controller.edit(
            "/Root",
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("ConstantString", "omni.graph.nodes.ConstantString"),
                ],
            },
        )
        controller = og.Controller(path_to_object_map=path_to_object_map, update_usd=False)
        (graph, _, _, _) = controller.edit(
            graph,
            {
                og.Controller.Keys.SET_VALUES: [
                    ("ConstantString.inputs:value", "foo"),
                ],
            },
        )
        string_attrib = controller.attribute("inputs:value", constant_string)
        string_usd_attrib = stage.GetAttributeAtPath(string_attrib.get_path())
        self.assertEqual("foo", controller.get(string_attrib))
        self.assertEqual(None, string_usd_attrib.Get())
        # flush to USD and verify the USD attrib value has changed
        bucket_id = constant_string.get_backing_bucket_id()
        graph.get_default_graph_context().write_bucket_to_backing(bucket_id)
        self.assertEqual("foo", string_usd_attrib.Get())

        # Verify changing OG attribute does not change USD attrib
        string_attrib.set("bar")
        await controller.evaluate()
        self.assertEqual("bar", controller.get(string_attrib))
        self.assertEqual("foo", string_usd_attrib.Get())
        # flush to USD and verify the USD attrib value has changed
        graph.get_default_graph_context().write_bucket_to_backing(bucket_id)
        self.assertEqual("bar", controller.get(string_attrib))
        self.assertEqual("bar", string_usd_attrib.Get())

        # verify invalid bucket raises
        # FC issues: getNamesAndTypes, bucketId 42 not found
        with ogts.ExpectedError():
            with self.assertRaises(ValueError):
                bad_bucket_id = og.BucketId(42)
                graph.get_default_graph_context().write_bucket_to_backing(bad_bucket_id)

        controller.edit("/Root", {og.Controller.Keys.DELETE_NODES: "ConstantString"})
        with self.assertRaises(ValueError):
            constant_string.get_backing_bucket_id()

    async def test_ogn_defaults(self):
        # tests that overriden attributes values persists after serialization
        # tests that non-overriden attributes values are set to OGN defaults, even after serialization

        stage = omni.usd.get_context().get_stage()
        (_, (data_node), _, _) = og.Controller.edit(
            "/World/Graph", {og.Controller.Keys.CREATE_NODES: [("TupleData", "omni.tutorials.TupleData")]}
        )
        self.assertTrue((data_node[0].get_attribute("inputs:a_double3").get() == [1.1, 2.2, 3.3]).all())

        # we don't have any value in USD
        node_prim = stage.GetPrimAtPath("/World/Graph/TupleData")
        usd_attrib = node_prim.GetAttribute("inputs:a_double3").Get()
        self.assertTrue(usd_attrib is None)

        # set runtime only
        data_node[0].get_attribute("inputs:a_double3").set([16.0, 17.0, 18.0])
        self.assertTrue((data_node[0].get_attribute("inputs:a_double3").get() == [16.0, 17.0, 18.0]).all())

        # does not affect USD
        usd_attrib = node_prim.GetAttribute("inputs:a_double3").Get()
        self.assertTrue(usd_attrib is None)

        # save and reload
        with tempfile.TemporaryDirectory() as test_directory:
            tmp_file_path = os.path.join(test_directory, "tmp_test_ogn_defaults.usda")
            (result, error, _) = await omni.usd.get_context().save_as_stage_async(str(tmp_file_path))
            self.assertTrue(result, error)
            (result, error) = await omni.usd.get_context().open_stage_async(tmp_file_path)
            self.assertTrue(result, error)

        stage = omni.usd.get_context().get_stage()
        node_prim = stage.GetPrimAtPath("/World/Graph/TupleData")

        # still no usd value
        usd_attrib = node_prim.GetAttribute("inputs:a_double3").Get()
        self.assertTrue(usd_attrib is None)

        # we have the default OGN values in OG
        attr = og.Controller.attribute("/World/Graph/TupleData/inputs:a_double3")
        self.assertTrue((attr.get() == [1.1, 2.2, 3.3]).all())

        # replace the node with something else in order to test a different default
        interface = ogu.get_node_type_forwarding_interface()
        try:
            self.assertTrue(
                interface.define_forward(
                    "omni.tutorials.TupleData", 1, "omni.tests.FakeTupleData", 1, "omni.graph.test"
                )
            )

            # save and reload
            with tempfile.TemporaryDirectory() as test_directory:
                tmp_file_path = os.path.join(test_directory, "tmp_test_ogn_defaults.usda")
                (result, error, _) = await omni.usd.get_context().save_as_stage_async(str(tmp_file_path))
                self.assertTrue(result, error)
                (result, error) = await omni.usd.get_context().open_stage_async(tmp_file_path)
                self.assertTrue(result, error)

            stage = omni.usd.get_context().get_stage()
            node_prim = stage.GetPrimAtPath("/World/Graph/TupleData")

            # still no usd value
            usd_attrib = node_prim.GetAttribute("inputs:a_double3").Get()
            self.assertTrue(usd_attrib is None)

            # we have now the new "fake"  default OGN values in OG
            attr = og.Controller.attribute("/World/Graph/TupleData/inputs:a_double3")
            self.assertTrue((attr.get() == [4.5, 6.7, 8.9]).all())

            # set USD
            node_prim.GetAttribute("inputs:a_double3").Set(Gf.Vec3f(6.0, 7.0, 8.0))
            self.assertTrue((attr.get() == [6.0, 7.0, 8.0]).all())

            # save and reload
            with tempfile.TemporaryDirectory() as test_directory:
                tmp_file_path = os.path.join(test_directory, "tmp_test_ogn_defaults.usda")
                (result, error, _) = await omni.usd.get_context().save_as_stage_async(str(tmp_file_path))
                self.assertTrue(result, error)
                (result, error) = await omni.usd.get_context().open_stage_async(tmp_file_path)
                self.assertTrue(result, error)

            # it should has now persisted
            stage = omni.usd.get_context().get_stage()
            node_prim = stage.GetPrimAtPath("/World/Graph/TupleData")
            self.assertTrue(node_prim.GetAttribute("inputs:a_double3").Get() == [6.0, 7.0, 8.0])

            # ...and proper OG values
            attr = og.Controller.attribute("/World/Graph/TupleData/inputs:a_double3")
            self.assertTrue((attr.get() == [6.0, 7.0, 8.0]).all())

        finally:
            interface.remove_forwarded_type("omni.tests.FakeTupleData", 1)


# ======================================================================
class TestOmniGraphSimpleSchema(ogts.OmniGraphTestCaseNoClear):
    """Test basic schema-based graph functionality"""

    # ----------------------------------------------------------------------
    async def test_node_duplication(self):
        """Test that duplication of an OmniGraph node copies the values correctly"""
        controller = og.Controller()
        (_, (node, py_node), _, _) = controller.edit(
            "/TestGraph",
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("SimpleData", "omni.graph.tutorials.SimpleData"),
                    ("SimpleDataPy", "omni.graph.tutorials.SimpleDataPy"),
                ],
                og.Controller.Keys.SET_VALUES: [
                    ("SimpleData.inputs:a_string", "Not The Default"),
                    ("SimpleData.inputs:a_constant_input", 17),
                    ("SimpleData.inputs:a_double", 3.25),
                    ("SimpleDataPy.inputs:a_string", "Not The Default"),
                    ("SimpleDataPy.inputs:a_constant_input", 17),
                    ("SimpleDataPy.inputs:a_double", 3.25),
                ],
            },
        )
        await controller.evaluate()

        # Confirm that the value being set made it to USD
        stage = omni.usd.get_context().get_stage()
        original_prim = stage.GetPrimAtPath("/TestGraph/SimpleData")
        value = original_prim.GetAttribute("inputs:a_string").Get()
        self.assertEqual(value, og.Controller(("inputs:a_string", node)).get())

        # Check that the original nodes have their values set properly
        self.assertEqual("Not The Default", og.Controller(("inputs:a_string", node)).get())
        self.assertEqual(17, og.Controller(("inputs:a_constant_input", node)).get())
        self.assertEqual(3.25, og.Controller(("inputs:a_double", node)).get())
        self.assertEqual(True, og.Controller(("inputs:a_bool", node)).get())
        self.assertEqual(0, og.Controller(("inputs:a_int", node)).get())

        self.assertEqual("Not The Default", og.Controller(("inputs:a_string", py_node)).get())
        self.assertEqual(17, og.Controller(("inputs:a_constant_input", py_node)).get())
        self.assertEqual(3.25, og.Controller(("inputs:a_double", py_node)).get())
        self.assertEqual(True, og.Controller(("inputs:a_bool", py_node)).get())
        self.assertEqual(0, og.Controller(("inputs:a_int", py_node)).get())

        # Duplicate the node
        omni.kit.commands.execute(
            "CopyPrims",
            paths_from=[str(node.get_prim_path()), str(py_node.get_prim_path())],
            paths_to=["/TestGraph/Duplicate", "/TestGraph/DuplicatePy"],
        )
        await omni.kit.app.get_app().next_update_async()

        # Check that the duplicate has its values set to the ones supplied by the original node
        duplicate_node = controller.node("/TestGraph/Duplicate")
        self.assertIsNotNone(duplicate_node)
        self.assertTrue(duplicate_node.is_valid())
        self.assertEqual("Not The Default", og.Controller(("inputs:a_string", duplicate_node)).get())
        self.assertEqual(17, og.Controller(("inputs:a_constant_input", duplicate_node)).get())
        self.assertEqual(3.25, og.Controller(("inputs:a_double", duplicate_node)).get())
        self.assertEqual(True, og.Controller(("inputs:a_bool", duplicate_node)).get())
        self.assertEqual(0, og.Controller(("inputs:a_int", duplicate_node)).get())

        duplicate_py_node = controller.node("/TestGraph/DuplicatePy")
        self.assertIsNotNone(duplicate_py_node)
        self.assertTrue(duplicate_py_node.is_valid())
        self.assertEqual("Not The Default", og.Controller(("inputs:a_string", duplicate_py_node)).get())
        self.assertEqual(17, og.Controller(("inputs:a_constant_input", duplicate_py_node)).get())
        self.assertEqual(3.25, og.Controller(("inputs:a_double", duplicate_py_node)).get())
        self.assertEqual(True, og.Controller(("inputs:a_bool", duplicate_py_node)).get())
        self.assertEqual(0, og.Controller(("inputs:a_int", duplicate_py_node)).get())
