# noqa: PLC0302
from math import isnan

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
from pxr import Sdf, Usd, UsdGeom


# Tests for ReadPrim (singular) node
class TestReadPrimNode(ogts.OmniGraphTestCase):
    """Unit tests for ReadPrim node in this extension"""

    async def test_read_prim_node_v1(self):
        """Validate backward compatibility for legacy versions of ReadPrim node."""
        # load the test scene which contains a ReadPrim V1 node
        (result, error) = await ogts.load_test_file("TestReadPrimNode_v1.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        test_graph_path = "/World/TestGraph"
        test_graph = og.get_graph_by_path(test_graph_path)
        read_prim_node = test_graph.get_node(test_graph_path + "/read_prim")
        self.assertTrue(read_prim_node.is_valid())

        # validate new attributes can be automatically created with their default value set.
        self.assertTrue(read_prim_node.get_attribute_exists("inputs:computeBoundingBox"))
        self.assertFalse(read_prim_node.get_attribute("inputs:computeBoundingBox").get())
        self.assertTrue(read_prim_node.get_attribute_exists("inputs:usdTimecode"))
        self.assertTrue(isnan(read_prim_node.get_attribute("inputs:usdTimecode").get()))


# Tests for ReadPrimBundle node
class TestReadPrimsBundle(ogts.OmniGraphTestCase):
    """Unit tests for ReadPrimBundle node in this extension"""

    async def test_read_prims_bundle_node_target_mpib(self):
        """Validate backward compatibility for legacy versions of ReadPrim node."""
        (result, error) = await ogts.load_test_file("TestReadPrimNode_targetMPiB.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        test_graph_path = "/TestGraph"
        test_graph = og.get_graph_by_path(test_graph_path)
        test_context = test_graph.get_default_graph_context()
        read_prim_bundle_node = test_graph.get_node(test_graph_path + "/read_prims_into_bundle")
        self.assertTrue(read_prim_bundle_node.is_valid())

        factory = og.IBundleFactory.create()
        output_bundle = test_context.get_output_bundle(read_prim_bundle_node, "outputs_primsBundle")
        output_bundle2 = factory.get_bundle(test_context, output_bundle)

        self.assertTrue(output_bundle2.valid)
        self.assertTrue(output_bundle2.get_child_bundle_count(), 2)

        child0 = output_bundle2.get_child_bundle(0)
        child1 = output_bundle2.get_child_bundle(1)
        self.assertTrue(child0.valid)
        self.assertTrue(child1.valid)

        self.assertTrue("size" in child0.get_attribute_names())
        self.assertTrue("size" in child1.get_attribute_names())

    async def test_read_prims_bundle_node_path_mpib(self):
        """Load multiple primitives into a bundle through input path"""
        (result, error) = await ogts.load_test_file("TestReadPrimNode_pathMPiB.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        test_graph_path = "/TestGraph"
        test_graph = og.get_graph_by_path(test_graph_path)
        test_context = test_graph.get_default_graph_context()
        read_prim_bundle_node = test_graph.get_node(test_graph_path + "/read_prims_into_bundle")
        self.assertTrue(read_prim_bundle_node.is_valid())

        factory = og.IBundleFactory.create()
        output_bundle = test_context.get_output_bundle(read_prim_bundle_node, "outputs_primsBundle")
        output_bundle2 = factory.get_bundle(test_context, output_bundle)

        self.assertTrue(output_bundle2.valid)
        self.assertTrue(output_bundle2.get_child_bundle_count(), 2)

        child0 = output_bundle2.get_child_bundle(0)
        child1 = output_bundle2.get_child_bundle(1)
        self.assertTrue(child0.valid)
        self.assertTrue(child1.valid)

        self.assertTrue("size" in child0.get_attribute_names())
        self.assertTrue("size" in child1.get_attribute_names())


# Tests for both ReadPrims (legacy) and ReadPrimsV2 node
# Must be inherited by a child test class to be tested
class TestReadPrimsCommon:
    """Unit tests for both ReadPrims (legacy) and ReadPrimsV2 node in this extension"""

    def __init__(self, read_prims_node_type):
        self._read_prims_node_type = read_prims_node_type

    async def test_import_matrix_and_bbox(self):
        """Verify that bounding box and matrix are correctly computed by ReadPrim."""
        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        # create 500 cubes to feed ReadPrims and trigger its concurrent collection of prims
        cube_paths = []
        cubes = []

        for i in range(500):
            cube_path = "/World/Cube" + str(i)
            cube = ogts.create_cube(stage, cube_path[1:], (1, 1, 1))
            UsdGeom.Xformable(cube).AddTranslateOp()

            cube_paths.append(cube_path)
            cubes.append(cube)

            self.assertEqual(cube.GetPath(), cube_path)  # noqa: PLE1101

        (_, nodes, _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("ReadPrims", self._read_prims_node_type),
                    ("ExtractPrim", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractBundle", "omni.graph.nodes.ExtractBundle"),
                    ("Inv", "omni.graph.nodes.OgnInvertMatrix"),
                    ("Const0", "omni.graph.nodes.ConstantDouble3"),
                    ("Const1", "omni.graph.nodes.ConstantDouble3"),
                ],
                keys.SET_VALUES: [
                    ("ExtractPrim.inputs:primPath", cube_paths[0]),
                ],
                keys.CONNECT: [
                    ("ReadPrims.outputs_primsBundle", "ExtractPrim.inputs:prims"),
                    ("ExtractPrim.outputs_primBundle", "ExtractBundle.inputs:bundle"),
                ],
            },
        )

        # add target prim
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath("/TestGraph/ReadPrims.inputs:prims"),
            targets=cube_paths,
        )

        # get initial output attributes
        await controller.evaluate()

        # connect the matrix
        controller.edit("/TestGraph", {keys.CONNECT: ("ExtractBundle.outputs:worldMatrix", "Inv.inputs:matrix")})

        await controller.evaluate()

        read_prims_node = nodes[0]
        extract_bundle = nodes[2]

        # We did not ask for BBox
        self.assertFalse(read_prims_node.get_attribute_exists("outputs:bboxMinCorner"))  # noqa: PLE1101
        self.assertFalse(read_prims_node.get_attribute_exists("outputs:bboxMaxCorner"))  # noqa: PLE1101
        self.assertFalse(read_prims_node.get_attribute_exists("outputs:bboxTransform"))  # noqa: PLE1101

        # but we should have the default matrix
        mat_attr = extract_bundle.get_attribute("outputs:worldMatrix")
        mat = mat_attr.get()
        self.assertTrue((mat == [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]).all())  # noqa: PLE1101

        # move the cube, check the new matrix
        pos = cubes[0].GetAttribute("xformOp:translate")
        pos.Set((2, 2, 2))

        await controller.evaluate()

        mat = mat_attr.get()
        self.assertTrue((mat == [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 2, 2, 2, 1]).all())  # noqa: PLE1101

        # now ask for bbox, and check its value
        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("ReadPrims.inputs:computeBoundingBox", True),
            },
        )
        await controller.evaluate()

        self.assertTrue(extract_bundle.get_attribute_exists("outputs:bboxMinCorner"))  # noqa: PLE1101
        self.assertTrue(extract_bundle.get_attribute_exists("outputs:bboxMaxCorner"))  # noqa: PLE1101
        self.assertTrue(extract_bundle.get_attribute_exists("outputs:bboxTransform"))  # noqa: PLE1101

        # connect the bbox
        controller.edit(
            "/TestGraph",
            {
                keys.CONNECT: [
                    ("ExtractBundle.outputs:bboxMinCorner", "Const0.inputs:value"),
                    ("ExtractBundle.outputs:bboxMaxCorner", "Const1.inputs:value"),
                ]
            },
        )

        await controller.evaluate()

        bbox_min = extract_bundle.get_attribute("outputs:bboxMinCorner")
        bbox_max = extract_bundle.get_attribute("outputs:bboxMaxCorner")

        bbox_min_val = bbox_min.get()
        bbox_max_val = bbox_max.get()

        self.assertTrue((bbox_min_val == [-0.5, -0.5, -0.5]).all())  # noqa: PLE1101
        self.assertTrue((bbox_max_val == [0.5, 0.5, 0.5]).all())  # noqa: PLE1101

        # then resize the cube, and check new bbox
        ext = cubes[0].GetAttribute("extent")
        ext.Set([(-10, -10, -10), (10, 10, 10)])

        await controller.evaluate()

        bbox_min_val = bbox_min.get()
        bbox_max_val = bbox_max.get()

        self.assertTrue((bbox_min_val == [-10, -10, -10]).all())  # noqa: PLE1101
        self.assertTrue((bbox_max_val == [10, 10, 10]).all())  # noqa: PLE1101

    async def _test_apply_skel_binding_on_skinned_model(self, test_scene):
        """shared function to validate applying skel binding on a model which has SkelBindingAPI schema."""
        # load test scene which contains a mesh that is bound with a skeleton and a skel animation
        (result, error) = await ogts.load_test_file(test_scene, use_caller_subdirectory=True)
        self.assertTrue(result, error)  # noqa: PLE1101

        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        bundle_factory = og.IBundleFactory.create()

        # create a ReadPrims node
        (test_graph, [read_prims_node], _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("read_prims", self._read_prims_node_type),
                ],
            },
        )

        # add target prim
        cylinder_path = "/Root/group1/pCylinder1_Skinned"
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath("/TestGraph/read_prims.inputs:prims"),
            target=cylinder_path,
        )

        await controller.evaluate()

        # validate a single child bundle is created when applySkelBinding is not enabled yet.
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        output_bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(output_bundle.valid)  # noqa: PLE1101

        child_bundles = output_bundle.get_child_bundles()
        self.assertTrue(len(child_bundles), 1)  # noqa: PLE1101

        single_prim_in_bundle = child_bundles[0]
        self.assertTrue(single_prim_in_bundle.valid)  # noqa: PLE1101

        source_prim_path_str = "sourcePrimPath"

        source_prim_path_attr = single_prim_in_bundle.get_attribute_by_name(source_prim_path_str)
        self.assertTrue(source_prim_path_attr.is_valid())  # noqa: PLE1101
        self.assertEqual(source_prim_path_attr.get(), cylinder_path)  # noqa: PLE1101

        attributes_without_skel_binding = set(single_prim_in_bundle.get_attribute_names())
        self.assertTrue("points" in attributes_without_skel_binding)  # noqa: PLE1101
        self.assertTrue("normals" in attributes_without_skel_binding)  # noqa: PLE1101

        # apply skel binding
        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("read_prims.inputs:applySkelBinding", True),
            },
        )
        await controller.evaluate()

        # validate child bundles are created for cylinder, skeleton and skel animation.
        child_bundles = output_bundle.get_child_bundles()
        self.assertEqual(len(child_bundles), 3)  # noqa: PLE1101

        cylinder_bundle = None
        skeleton_bundle = None
        skelanim_bundle = None

        for child_bundle in child_bundles:
            source_prim_path_attr = child_bundle.get_attribute_by_name(source_prim_path_str)
            if source_prim_path_attr.get() == cylinder_path:
                cylinder_bundle = child_bundle
            else:
                source_prim_type_attr = child_bundle.get_attribute_by_name("sourcePrimType")
                source_prim_type = source_prim_type_attr.get()
                if source_prim_type == "Skeleton":
                    skeleton_bundle = child_bundle
                elif source_prim_type == "SkelAnimation":
                    skelanim_bundle = child_bundle

        self.assertIsNotNone(cylinder_bundle)  # noqa: PLE1101
        self.assertIsNotNone(skeleton_bundle)  # noqa: PLE1101
        self.assertIsNotNone(skelanim_bundle)  # noqa: PLE1101

        # the cylinder bundle has the skel:skeleton attribute that stores the link to the Skeleton bundle
        skeleton_path_attr = cylinder_bundle.get_attribute_by_name("skel:skeleton")
        self.assertTrue(skeleton_path_attr.is_valid())  # noqa: PLE1101

        source_prim_path_attr = skeleton_bundle.get_attribute_by_name(source_prim_path_str)
        self.assertTrue(source_prim_path_attr.is_valid())  # noqa: PLE1101

        skeleton_path = skeleton_path_attr.get()
        self.assertEqual(source_prim_path_attr.get(), skeleton_path)  # noqa: PLE1101

        # the skeleton bundle has the skel:animationSource attribute that stores the link to the SkelAnimation bundle
        skelanim_path_attr = skeleton_bundle.get_attribute_by_name("skel:animationSource")
        self.assertTrue(skelanim_path_attr.is_valid())  # noqa: PLE1101

        source_prim_path_attr = skelanim_bundle.get_attribute_by_name(source_prim_path_str)
        self.assertTrue(source_prim_path_attr.is_valid())  # noqa: PLE1101

        skelanim_path = skelanim_path_attr.get()
        self.assertEqual(source_prim_path_attr.get(), skelanim_path)  # noqa: PLE1101

        # validate additional skel attributes added to the cylinder bundle
        attributes_with_skel_binding = set(cylinder_bundle.get_attribute_names())
        self.assertTrue(attributes_without_skel_binding.issubset(attributes_with_skel_binding))  # noqa: PLE1101

        additional_skel_attributes = attributes_with_skel_binding.difference(attributes_without_skel_binding)
        self.assertTrue("points:default" in additional_skel_attributes)  # noqa: PLE1101
        self.assertTrue("normals:default" in additional_skel_attributes)  # noqa: PLE1101

        # validate the joints attribute
        joints_attr = skeleton_bundle.get_attribute_by_name("joints")
        self.assertTrue(joints_attr.is_valid())  # noqa: PLE1101
        num_joints = joints_attr.size()

        # validate the rest pose
        matrix_attr_type = og.Type(og.BaseDataType.DOUBLE, 16, 1, og.AttributeRole.MATRIX)
        rest_transforms_attr = skeleton_bundle.get_attribute_by_name("restTransforms")
        self.assertTrue(rest_transforms_attr.is_valid())  # noqa: PLE1101
        self.assertEqual(rest_transforms_attr.get_type(), matrix_attr_type)  # noqa: PLE1101
        self.assertEqual(rest_transforms_attr.size(), num_joints)  # noqa: PLE1101

        # validate the animated pose which is stored at the "jointLocalTransforms" attribute
        joint_local_transforms_attr = skeleton_bundle.get_attribute_by_name("jointLocalTransforms")
        self.assertTrue(joint_local_transforms_attr.is_valid())  # noqa: PLE1101
        self.assertEqual(joint_local_transforms_attr.get_type(), matrix_attr_type)  # noqa: PLE1101
        self.assertEqual(joint_local_transforms_attr.size(), num_joints)  # noqa: PLE1101

    async def test_apply_skel_binding_on_authored_skeletal_relationships(self):
        """validate applying skel binding on a prim that has SkelBindingAPI schema and authored skeletal relationships."""
        await self._test_apply_skel_binding_on_skinned_model("skelcylinder.usda")

    async def test_apply_skel_binding_on_inherited_skeletal_relationships(self):
        """validate applying skel binding on a prim that has SkelBindingAPI schema and inherited skeletal relationships."""
        await self._test_apply_skel_binding_on_skinned_model("skelcylinder_inherited_skeleton.usda")

    async def test_apply_skel_binding_on_model_without_normals(self):
        """validate applying skel binding on a model which has SkelBindingAPI schema and no normals."""
        # load test scene which contains a mesh that is bound with a skeleton and a skel animation
        (result, error) = await ogts.load_test_file("skel_model_without_normals.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)  # noqa: PLE1101

        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        bundle_factory = og.IBundleFactory.create()

        # create a ReadPrims node
        (test_graph, [read_prims_node], _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("read_prims", self._read_prims_node_type),
                ],
                keys.SET_VALUES: [
                    # Remove the normals attribute from the list of attributes to import, because ReadPrims
                    # imports normals by default, and we want to test the SkelBinding when the normals are not
                    # imported.
                    ("read_prims.inputs:attrNamesToImport", "* ^normals")
                ],
            },
        )

        # add target prim
        mesh_path = "/Model/Arm"
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath("/TestGraph/read_prims.inputs:prims"),
            target=mesh_path,
        )

        await controller.evaluate()

        # validate a single child bundle is created when applySkelBinding is not enabled yet.
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        output_bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(output_bundle.valid)  # noqa: PLE1101

        child_bundles = output_bundle.get_child_bundles()
        self.assertTrue(len(child_bundles), 1)  # noqa: PLE1101

        mesh_bundle = child_bundles[0]
        self.assertTrue(mesh_bundle.valid)  # noqa: PLE1101

        # the mesh bundle has points but no normals
        attributes_without_skel_binding = set(mesh_bundle.get_attribute_names())
        self.assertTrue("points" in attributes_without_skel_binding)  # noqa: PLE1101
        self.assertFalse("normals" in attributes_without_skel_binding)  # noqa: PLE1101

        # apply skel binding
        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("read_prims.inputs:applySkelBinding", True),
            },
        )
        await controller.evaluate()

        # there will be three bundles and we need to get the mesh bundles
        child_bundles = output_bundle.get_child_bundles()
        self.assertEqual(len(child_bundles), 3)  # noqa: PLE1101

        mesh_bundle = None

        for child_bundle in child_bundles:
            source_prim_path_attr = child_bundle.get_attribute_by_name("sourcePrimPath")
            if source_prim_path_attr.get() == mesh_path:
                mesh_bundle = child_bundle
                break

        self.assertIsNotNone(mesh_bundle)  # noqa: PLE1101

        # validate additional skel attributes added to the cylinder bundle
        attributes_with_skel_binding = set(mesh_bundle.get_attribute_names())
        self.assertTrue(attributes_without_skel_binding.issubset(attributes_with_skel_binding))  # noqa: PLE1101

        additional_skel_attributes = attributes_with_skel_binding.difference(attributes_without_skel_binding)
        self.assertTrue("points:default" in additional_skel_attributes)  # noqa: PLE1101
        self.assertFalse("normals:default" in additional_skel_attributes)  # noqa: PLE1101

    async def test_apply_skel_binding_on_prim_without_required_schema(self):
        """validate applying skel binding to a simple cube which has no SkelBindingAPI schema."""
        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        bundle_factory = og.IBundleFactory.create()

        # create a ReadPrims node
        (test_graph, [read_prims_node], _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("read_prims", self._read_prims_node_type),
                ],
            },
        )

        # add target prim
        cube = ogts.create_cube(stage, "World/Cube", (1, 1, 1))
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath("/TestGraph/read_prims.inputs:prims"),
            target=cube.GetPath(),
        )

        await controller.evaluate()

        # validate a single child bundle is created before applying skel binding.
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        output_bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(output_bundle.valid)  # noqa: PLE1101

        child_bundles = output_bundle.get_child_bundles()
        self.assertTrue(len(child_bundles), 1)  # noqa: PLE1101

        child_bundle_before_skel_binding = child_bundles[0]
        self.assertTrue(child_bundle_before_skel_binding.valid)  # noqa: PLE1101

        source_prim_path_str = "sourcePrimPath"

        source_prim_path_attr = child_bundle_before_skel_binding.get_attribute_by_name(source_prim_path_str)
        self.assertTrue(source_prim_path_attr.is_valid())  # noqa: PLE1101
        self.assertEqual(source_prim_path_attr.get(), cube.GetPath())  # noqa: PLE1101

        attributes_before_skel_binding = set(child_bundle_before_skel_binding.get_attribute_names())

        # apply skel binding
        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("read_prims.inputs:applySkelBinding", True),
            },
        )
        await controller.evaluate()

        # validate no additional child bundle is created after applying skel binding.
        child_bundles = output_bundle.get_child_bundles()
        self.assertEqual(len(child_bundles), 1)  # noqa: PLE1101

        child_bundle_after_skel_binding = child_bundles[0]
        self.assertTrue(child_bundle_after_skel_binding.valid)  # noqa: PLE1101

        source_prim_path_attr = child_bundle_after_skel_binding.get_attribute_by_name(source_prim_path_str)
        self.assertTrue(source_prim_path_attr.is_valid())  # noqa: PLE1101
        self.assertEqual(source_prim_path_attr.get(), cube.GetPath())  # noqa: PLE1101

        # validate no additional attributes have been added to the bundle
        attributes_after_skel_binding = set(child_bundle_after_skel_binding.get_attribute_names())
        self.assertEqual(attributes_before_skel_binding, attributes_after_skel_binding)  # noqa: PLE1101

    async def _test_apply_skel_binding_on_joint_ordering(self, test_scene, path_skinned_mesh, path_baked_mesh):
        """shared function to validate applying skel binding on joint ordering."""
        test_graph_path = "/TestGraph"
        time_code_range = range(1, 30)

        # load test scene which contains a mesh that is bound with a skeleton and a skel animation
        (result, error) = await ogts.load_test_file(test_scene, use_caller_subdirectory=True)
        self.assertTrue(result, error)  # noqa: PLE1101

        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        bundle_factory = og.IBundleFactory.create()

        (test_graph, (_, extracter_0, extracter_1), _, _) = controller.edit(
            test_graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadPrims", self._read_prims_node_type),
                    ("ExtractPrim_SkelBinding", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractPrim_BakeSkinning", "omni.graph.nodes.ExtractPrim"),
                ],
                keys.SET_VALUES: [
                    ("ReadPrims.inputs:applySkelBinding", True),
                    ("ExtractPrim_SkelBinding.inputs:primPath", path_skinned_mesh),
                    ("ExtractPrim_BakeSkinning.inputs:primPath", path_baked_mesh),
                ],
                keys.CONNECT: [
                    ("ReadPrims.outputs_primsBundle", "ExtractPrim_SkelBinding.inputs:prims"),
                    ("ReadPrims.outputs_primsBundle", "ExtractPrim_BakeSkinning.inputs:prims"),
                ],
            },
        )

        # Assign the target prims to the import node.
        cylinder_paths = [path_skinned_mesh, path_baked_mesh]
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(test_graph_path + "/ReadPrims.inputs:prims"),
            targets=cylinder_paths,
        )

        # get initial output attributes
        await controller.evaluate()

        # fetch skinned points
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(extracter_0, "outputs_primBundle")
        bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(bundle.valid)  # noqa: PLE1101
        skinned_points_attr = bundle.get_attribute_by_name("points")
        self.assertTrue(skinned_points_attr.is_valid())  # noqa: PLE1101

        # fetch baked points
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(extracter_1, "outputs_primBundle")
        bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(bundle.valid)  # noqa: PLE1101
        baked_points_attr = bundle.get_attribute_by_name("points")
        self.assertTrue(baked_points_attr.is_valid())  # noqa: PLE1101

        # playback and compare the two lists of points
        for time_code in time_code_range:
            controller.edit(
                test_graph_path,
                {keys.SET_VALUES: ("ReadPrims.inputs:usdTimecode", time_code)},
            )

            await og.Controller.evaluate(test_graph)

            skinned_points = skinned_points_attr.get().tolist()
            self.assertTrue(isinstance(skinned_points, list))  # noqa: PLE1101
            self.assertTrue(len(skinned_points) != 0)  # noqa: PLE1101

            baked_points = baked_points_attr.get().tolist()
            self.assertTrue(isinstance(baked_points, list))  # noqa: PLE1101
            self.assertTrue(len(baked_points) != 0)  # noqa: PLE1101

            self.assertTrue(skinned_points == baked_points)  # noqa: PLE1101

    async def test_apply_skel_binding_on_joint_ordering(self):
        """validate applying skel binding on joint ordering which is defined on skeleton but may or may not be overridden by skinned mesh"""
        # The skinned meshes are baked offline with the following script. Then we can compare the computation result of the
        # `ReadPrims` node with the baked meshes.
        #
        # from pxr import Gf, Usd, UsdSkel
        #
        # testFile = "lbs.usda"
        # stage = Usd.Stage.Open(testFile)
        #
        # UsdSkel.BakeSkinning(stage.Traverse())
        #
        # stage.GetRootLayer().Export("lbs.baked.usda")

        # Each tuple is (test_scene, path_skinned_mesh, path_baked_mesh)
        test_cases = [
            ("lbs.usda", "/Root/Mesh_Skinned", "/Root/Mesh_Baked"),
            ("lbs_custom_joint_ordering.usda", "/Root/Mesh_Skinned", "/Root/Mesh_Baked"),
            ("skelcylinder.usda", "/Root/group1/pCylinder1_Skinned", "/Root/group1/pCylinder1_Baked"),
            (
                "skelcylinder_custom_joint_ordering.usda",
                "/Root/group1/pCylinder1_Skinned",
                "/Root/group1/pCylinder1_Baked",
            ),
        ]

        for test_case in test_cases:
            await self._test_apply_skel_binding_on_joint_ordering(*test_case)

    async def test_read_prims_attribute_filter(self):
        """Validate the attribute filter functionality of ReadPrim."""
        await self._test_attribute_filter(self._read_prims_node_type)

    async def _test_attribute_filter(self, node_type):
        """Shared function to validate the attribute filter functionality of ReadPrimBundle and ReadPrim."""
        controller = og.Controller()
        keys = og.Controller.Keys

        test_graph_path = "/TestGraph"

        (test_graph, [read_prims_node], _, _) = controller.edit(
            test_graph_path,
            {
                keys.CREATE_NODES: [("read_prims_node", node_type)],
            },
        )

        # Create a cube and connect it as the input prim.
        stage = omni.usd.get_context().get_stage()
        cube = ogts.create_cube(stage, "World/Cube", (1, 1, 1))
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(test_graph_path + "/read_prims_node.inputs:prims"),
            target=cube.GetPath(),
        )
        await controller.evaluate()

        # Get the output bundle.
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        bundle_factory = og.IBundleFactory.create()
        output_bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(output_bundle.valid)  # noqa: PLE1101
        self.assertEqual(output_bundle.get_child_bundle_count(), 1)  # noqa: PLE1101

        # Get the only child bundle.
        child_bundle = output_bundle.get_child_bundle(0)
        self.assertTrue(child_bundle.valid)  # noqa: PLE1101

        # The "size" attribute is expected for a cube by default.
        default_attribute_names = child_bundle.get_attribute_names()
        self.assertTrue("size" in default_attribute_names)  # noqa: PLE1101

        # The filtered attribute names will be the same as the default.
        attr_filters = ["*", "* dummy", "dummy *"]

        for attr_filter in attr_filters:
            controller.edit(
                test_graph_path,
                {
                    keys.SET_VALUES: ("read_prims_node.inputs:attrNamesToImport", attr_filter),
                },
            )
            await controller.evaluate()

            self.assertEqual(set(child_bundle.get_attribute_names()), set(default_attribute_names))  # noqa: PLE1101

        # The filtered attribute names will contain "size" only.
        attr_filters = [
            "size",
            "size dummy",
            "dummy size",
            "????",
            "s???",
            "?ize",
            "?i?e",
            "*z?",
            "?im* ?ix* ?iz*",
            "???? dummy",
            "dummy ????",
        ]

        # "sourcePrimPath", "sourcePrimType" and "worldMatrix" are additional attributes created by the node.
        expected_attribute_names = {"sourcePrimPath", "sourcePrimType", "worldMatrix", "size"}

        for attr_filter in attr_filters:
            controller.edit(
                test_graph_path,
                {
                    keys.SET_VALUES: ("read_prims_node.inputs:attrNamesToImport", attr_filter),
                },
            )
            await controller.evaluate()

            self.assertEqual(set(child_bundle.get_attribute_names()), expected_attribute_names)  # noqa: PLE1101

        # The filtered attribute names will not contain "size".
        attr_filters = [
            " ",
            "Size",
            "SIZE",
            "dummy",
            "sizedummy",
            "dummysize",
            "size:dummy",
            "dummy:size",
            "size&dummy",
            "dummy&size",
            "*:dummy",
            "dummy:*",
            "???",
            "??",
            "?",
            "????:dummy",
            "dummy:????",
        ]

        # "sourcePrimPath", "sourcePrimType" and "worldMatrix" are additional attributes created by the node.
        expected_attribute_names = {"sourcePrimPath", "sourcePrimType", "worldMatrix"}

        for attr_filter in attr_filters:
            controller.edit(
                test_graph_path,
                {
                    keys.SET_VALUES: ("read_prims_node.inputs:attrNamesToImport", attr_filter),
                },
            )
            await controller.evaluate()

            self.assertEqual(set(child_bundle.get_attribute_names()), expected_attribute_names)  # noqa: PLE1101

    async def test_read_prims_time_code_mpib(self):
        """Test whether points and normals keep consistent after modifying time code."""
        controller = og.Controller()
        keys = og.Controller.Keys
        bundle_factory = og.IBundleFactory.create()
        stage = omni.usd.get_context().get_stage()

        # Create two cubes, whose types are Mesh and Cube respectively.
        cube_paths = ["/cube_mesh", "/cube_shape"]

        _, default_cube_mesh_path = omni.kit.commands.execute("CreateMeshPrim", prim_type="Cube")
        omni.kit.commands.execute("MovePrim", path_from=default_cube_mesh_path, path_to=cube_paths[0])

        og.Controller.create_prim(cube_paths[1], {}, "Cube")

        # Create a graph that reads the two prims and extract the bundle of the cube mesh
        test_graph_path = "/TestGraph"
        (test_graph, [_, extract_prim_node], _, _,) = controller.edit(
            test_graph_path,
            {
                keys.CREATE_NODES: [
                    ("read_prims", self._read_prims_node_type),
                    ("extract_prim", "omni.graph.nodes.ExtractPrim"),
                ],
                keys.SET_VALUES: [
                    ("extract_prim.inputs:primPath", cube_paths[0]),
                ],
                keys.CONNECT: [
                    ("read_prims.outputs_primsBundle", "extract_prim.inputs:prims"),
                ],
            },
        )

        # Assign the target prims to the import node.
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(test_graph_path + "/read_prims.inputs:prims"),
            targets=cube_paths,
        )

        await og.Controller.evaluate(test_graph)

        # Get output bundle.
        test_context = test_graph.get_default_graph_context()
        output_bundle = test_context.get_output_bundle(extract_prim_node, "outputs_primBundle")
        bundle = bundle_factory.get_bundle(test_context, output_bundle)
        self.assertTrue(bundle.valid)  # noqa: PLE1101

        # Get points
        points_attr = bundle.get_attribute_by_name("points")
        self.assertTrue(points_attr.is_valid())  # noqa: PLE1101
        points = points_attr.get().tolist()
        self.assertTrue(isinstance(points, list))  # noqa: PLE1101
        self.assertTrue(len(points) != 0)  # noqa: PLE1101

        # Get normals
        normals_attr = bundle.get_attribute_by_name("normals")
        self.assertTrue(normals_attr.is_valid())  # noqa: PLE1101
        normals = normals_attr.get().tolist()
        self.assertTrue(isinstance(normals, list))  # noqa: PLE1101
        self.assertTrue(len(normals) != 0)  # noqa: PLE1101

        # Modify timecode and trigger evaluation. Run 60 times to eliminate random factor.
        for timecode in range(60):
            controller.edit(
                test_graph_path,
                {keys.SET_VALUES: ("read_prims.inputs:usdTimecode", timecode)},
            )

            await og.Controller.evaluate(test_graph)

            # Test whether the points data keep consistent
            points_attr = bundle.get_attribute_by_name("points")
            self.assertTrue(points_attr.is_valid())  # noqa: PLE1101
            self.assertSequenceEqual(points_attr.get().tolist(), points)  # noqa: PLE1101

            # Test whether the points data keep consistent
            normals_attr = bundle.get_attribute_by_name("normals")
            self.assertTrue(normals_attr.is_valid())  # noqa: PLE1101
            self.assertSequenceEqual(normals_attr.get().tolist(), normals)  # noqa: PLE1101

    async def test_read_prims_attribute_interpolation(self):
        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()

        stage = usd_context.get_stage()

        # Create Mesh and Primvar with interpolation
        mesh = ogts.create_grid_mesh(stage, "/Mesh")
        sts_primvar = UsdGeom.PrimvarsAPI(mesh).CreatePrimvar("st", Sdf.ValueTypeNames.Float2Array)
        sts_primvar.SetInterpolation("faceVarying")

        # Create Read Prim with test node to test if interpolation was read
        test_graph_path = "/TestGraph"
        (test_graph, [_, attr_interp], _, _,) = controller.edit(
            test_graph_path,
            {
                keys.CREATE_NODES: [
                    ("read_prims", self._read_prims_node_type),
                    ("attr_interp", "omni.graph.test.TestBundleAttributeInterpolation"),
                ],
                keys.SET_VALUES: ("attr_interp.inputs:attribute", "primvars:st"),
                keys.CONNECT: ("read_prims.outputs_primsBundle", "attr_interp.inputs:prims"),
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(test_graph_path + "/read_prims.inputs:prims"),
            target=mesh.GetPath(),
        )

        await og.Controller.evaluate(test_graph)

        # Confirm the interpolation is carried properly from Read Prims and Bundle Prims
        self.assertEqual(attr_interp.get_attribute("outputs:interpolation").get(), "faceVarying")  # noqa: PLE1101

    async def test_read_prim_nodes_have_nan_defaults(self):
        """Validates that manually created prim nodes have nan as default"""
        controller = og.Controller()
        keys = og.Controller.Keys
        test_graph_path = "/TestGraph"
        (_, nodes, _, _) = controller.edit(
            test_graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadPrims", self._read_prims_node_type),
                    ("ReadPrimAttr", "omni.graph.nodes.ReadPrimAttribute"),
                    ("ReadPrimAttrs", "omni.graph.nodes.ReadPrimAttributes"),
                    ("ReadPrimsBundle", "omni.graph.nodes.ReadPrimsBundle"),
                ],
            },
        )

        for node in nodes:
            self.assertTrue(isnan(node.get_attribute("inputs:usdTimecode").get()))  # noqa: PLE1101

    async def test_read_prim_have_no_errors_on_creation(self):
        """Validates that creating ReadPrimAttribute nodes with default values doesn't throw errors"""
        controller = og.Controller()
        keys = og.Controller.Keys
        test_graph_path = "/TestGraph"
        (test_graph, nodes, _, _) = controller.edit(
            {"graph_path": test_graph_path},
            {
                keys.CREATE_NODES: [
                    ("ReadPrimAttr", "omni.graph.nodes.ReadPrimAttribute"),
                ],
                keys.SET_VALUES: [
                    ("ReadPrimAttr.inputs:name", ""),
                ],
            },
        )

        self.assertEqual(len(nodes[0].get_compute_messages(og.ERROR)), 0)  # noqa: PLE1101
        nodes[0].clear_old_compute_messages()

        with ogts.ExpectedError():
            await og.Controller.evaluate(test_graph)
            self.assertEqual(len(nodes[0].get_compute_messages(og.ERROR)), 1)  # noqa: PLE1101


# Tests for ReadPrims (legacy) node only + all tests in TestReadPrimsCommon
class TestReadPrimsNode(ogts.OmniGraphTestCase, TestReadPrimsCommon):
    """
    Unit tests for ReadPrims (legacy) node in this extension
    It also inherits all tests from TestReadPrimsCommon unless explicitly opted out
    """

    def __init__(self, *args, **kwargs):
        ogts.OmniGraphTestCase.__init__(self, *args, **kwargs)
        TestReadPrimsCommon.__init__(self, read_prims_node_type="omni.graph.nodes.ReadPrims")

    async def test_read_prims_node_v1(self):
        """Validate backward compatibility for legacy versions of ReadPrims node."""
        # load the test scene which contains a ReadPrims V1 node
        (result, error) = await ogts.load_test_file("TestReadPrimsNode_v1.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        test_graph_path = "/World/TestGraph"
        test_graph = og.get_graph_by_path(test_graph_path)
        read_prim_node = test_graph.get_node(test_graph_path + "/read_prims")
        self.assertTrue(read_prim_node.is_valid())

        # validate new attributes can be automatically created with their default value set.
        attr_names_and_values = [
            ("inputs:applySkelBinding", False),
        ]

        for attr_name, attr_value in attr_names_and_values:
            self.assertTrue(read_prim_node.get_attribute_exists(attr_name))
            self.assertTrue(read_prim_node.get_attribute(attr_name).get() == attr_value)

    # this is a test for ReadPrimBundle, to utilize _test_attribute_filter, thus it is not in TestReadPrimsBundle
    async def test_read_prims_bundle_attribute_filter(self):
        """Validate the attribute filter functionality of ReadPrimBundle."""
        await self._test_attribute_filter("omni.graph.nodes.ReadPrimsBundle")


# Tests for ReadPrimsV2 node only + all tests in TestReadPrimsCommon
class TestReadPrimsV2Node(ogts.OmniGraphTestCase, TestReadPrimsCommon):
    """
    Unit tests for ReadPrimsV2 node in this extension
    It also inherits all tests from TestReadPrimsCommon unless explicitly opted out
    """

    def __init__(self, *args, **kwargs):
        ogts.OmniGraphTestCase.__init__(self, *args, **kwargs)
        TestReadPrimsCommon.__init__(self, read_prims_node_type="omni.graph.nodes.ReadPrimsV2")

    async def test_read_prims_root_prims_targets(self):
        """Test omni.graph.nodes.ReadPrims with rootPrims targets"""
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()
        test_graph_path = "/World/TestGraph"

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

        root1 = stage.DefinePrim("/Root1", "Xform")
        stage.DefinePrim("/Root1/Xform", "Xform")
        stage.DefinePrim("/Root1/Xform/Cube", "Cube")

        stage.DefinePrim("/Root2", "Xform")
        stage.DefinePrim("/Root2/Xform", "Xform")
        stage.DefinePrim("/Root2/Xform/Sphere", "Sphere")

        root3 = stage.DefinePrim("/Root3", "Xform")
        stage.DefinePrim("/Root3/Xform", "Xform")
        stage.DefinePrim("/Root3/Xform/Cone", "Cone")

        (graph, [read_prims_node], _, _) = controller.edit(
            test_graph_path,
            {
                keys.CREATE_NODES: [
                    ("Read", self._read_prims_node_type),
                ],
                keys.SET_VALUES: [
                    ("Read.inputs:pathPattern", "*"),
                ],
            },
        )

        # connect /Root1 and /Root3 to inputs:rootPrims
        # The node should find all prims under them, but not /Root2
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{test_graph_path}/Read.inputs:prims"),
            targets=[root1.GetPath(), root3.GetPath()],
        )

        await controller.evaluate()

        graph_context = graph.get_default_graph_context()

        output_prims_bundle = graph_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        self.assertTrue(output_prims_bundle.valid)

        expected_matched_prim_paths = {
            "/Root1",
            "/Root1/Xform",
            "/Root1/Xform/Cube",
            "/Root3",
            "/Root3/Xform",
            "/Root3/Xform/Cone",
        }

        child_bundles = output_prims_bundle.get_child_bundles()
        self.assertEqual(len(child_bundles), len(expected_matched_prim_paths))

        for bundle in child_bundles:
            path = bundle.get_attribute_by_name("sourcePrimPath").get()
            expected_matched_prim_paths.remove(path)

        # no exception should raise when remove a path in the loop above
        # after the loop, all paths should have been removed
        self.assertEqual(len(expected_matched_prim_paths), 0)
