## Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
##
## NVIDIA CORPORATION and its licensors retain all intellectual property
## and proprietary rights in and to this software, related documentation
## and any modifications thereto.  Any use, reproduction, disclosure or
## distribution of this software and related documentation without an express
## license agreement from NVIDIA CORPORATION is strictly prohibited.
##

#TODO 02: omnihydra does not support dynamically switch animationsource on skelroot prim bindingAPI (should be a repopulate in this case)
#TODO 03: omnihydra does not support joints/blendshapes change in skelanimation prim
#TODO 04: omnihydra does not support blendshapes (target name) change in skelanimation prim (not in test)
#TODO 05 when skelanimation become invalid, skinning result should retrive restTransform
#TODO 06: after added time ranged control, the render result does not work correct for the specific animation source

import omni.kit.app
import omni.kit.commands
import omni.kit.undo
import omni.kit.test
import omni.timeline
import carb.settings
from .test_hydra_common import RtxHydraTest
from .test_common import testSettings, postLoadSkelTestSettings, set_transform_helper, wait_for_update
from pxr import Gf, Sdf, Usd, UsdGeom, UsdSkel, UsdShade

def _update_animation(animation : UsdSkel.Animation, animation_static : UsdSkel.Animation, timecode : Usd.TimeCode):
    trans = animation.GetTranslationsAttr().Get(timecode)
    rots = animation.GetRotationsAttr().Get(timecode)
    scales = animation.GetScalesAttr().Get(timecode)
    bsWeights = animation.GetBlendShapeWeightsAttr().Get(timecode)
    animation_static.GetTranslationsAttr().Set(trans)
    animation_static.GetRotationsAttr().Set(rots)
    animation_static.GetScalesAttr().Set(scales)
    animation_static.GetBlendShapeWeightsAttr().Set(bsWeights)

class TestRtxHydraSkel(RtxHydraTest):
    TEST_PATH = "hydra/skel"

    async def setUp(self):
        await super().setUp()
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_fast_mode(True)
        await omni.kit.app.get_app().next_update_async()

    async def test_01_skel_anim(self):
        """
        Test hydra skel - skel anim
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skelroot_path = "/Root/group1"
        skelroot_prim = stage.GetPrimAtPath(skelroot_path)
        skelroot = UsdSkel.Root(skelroot_prim)
        skeleton_path = "/Root/group1/joint1"
        skeleton_prim = stage.GetPrimAtPath(skeleton_path)
        skeleton_bindingAPI = UsdSkel.BindingAPI(skeleton_prim)
        skeleton = UsdSkel.Skeleton(skeleton_prim)
        animation_static_path = "/Root/group1/joint1/Animation_Static"
        animation_static_prim = stage.GetPrimAtPath(animation_static_path)
        animation_static = UsdSkel.Animation(animation_static_prim)
        animation_path = "/Root/group1/joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        animation = UsdSkel.Animation(animation_prim)
        animation_flat_path = "/Root/group1/joint1/Animation_Flat"
        animation_flat_prim = stage.GetPrimAtPath(animation_flat_path)
        aniamtion_flat = UsdSkel.Animation(animation_flat_prim)
        animation_outside_path = "/ZAnimation"
        animation_outside_prim = stage.GetPrimAtPath(animation_outside_path)
        aniamtion_outside = UsdSkel.Animation(animation_outside_prim)

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(0.0)
        await omni.kit.app.get_app().next_update_async()

        # pure skeleton test
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_0_skeleton_0.png")

        #Test animation
        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_path])
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
        timeline.set_current_time(0.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_0.png")
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_1.png")
        timeline.set_current_time(1.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_2.png")
        skeleton_bindingAPI.GetAnimationSourceRel().ClearTargets(False)
        self.assertTrue(not skeleton_bindingAPI.GetInheritedAnimationSource())
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_3.png")

        #Test animation in session layer (resync)
        with Usd.EditContext(stage, session_layer):
            skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_path])
            self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
            timeline.set_current_time(0.0)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_sessionlayer_0.png")
            timeline.set_current_time(0.5)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_sessionlayer_1.png")
            timeline.set_current_time(1.0)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_sessionlayer_2.png")
            stage.RemovePrim(skeleton_path)
            self.assertTrue(not skeleton_bindingAPI.GetInheritedAnimationSource())
            timeline.set_current_time(0.5)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_1_animation_sessionlayer_3.png")

        await wait_for_update()
        self.assertTrue(not skeleton_bindingAPI.GetInheritedAnimationSource())
        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_path])
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)

        #Test animation switch
        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_flat_path])
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_flat_prim)
        timeline.set_current_time(0.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_2_animation_flat_0.png")
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_2_animation_flat_1.png")

        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_path])
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_2_animation_flat_2.png")

        #Test animation switch in session layer
        with Usd.EditContext(stage, session_layer):
            skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_flat_path])
            timeline.set_current_time(0.0)
            self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_flat_prim)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_2_animation_flat_sessionlayer_0.png")
            timeline.set_current_time(0.5)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_2_animation_flat_sessionlayer_1.png")
            #TODO 06: after added time ranged control, the render result does not work correct for the specific animation source
            stage.RemovePrim(skeleton_path)
            self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
            await wait_for_update()
            #await self.capture_and_compare(self.TEST_PATH, "01_skelanim_2_animation_flat_sessionlayer_2.png")

        #Test animation switch to outside root animation
        await wait_for_update()
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_outside_path])
        await wait_for_update()
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_outside_prim)
        timeline.set_current_time(0.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_3_animation_outside_0.png")
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_skelanim_3_animation_outside_1.png")

        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_path])
        await wait_for_update()
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)

        #Test animation switch to outside root animation in session layer
        with Usd.EditContext(stage, session_layer):
            skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_outside_path])
            await wait_for_update()
            self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_outside_prim)
            timeline.set_current_time(0.0)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_3_animation_outside_sessionlayer_0.png")
            timeline.set_current_time(0.5)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "01_skelanim_3_animation_outside_sessionlayer_1.png")

            stage.RemovePrim(skeleton_path)

        ##TODO 02: omnihydra does not support dynamically switch animationsource on skelroot prim bindingAPI (should be a repopulate in this case)
        ##Test animation switch to at skelroot prim
        #skeleton_bindingAPI.GetAnimationSourceRel().ClearTargets(False)
        #await wait_for_update()
        #self.assertTrue(not skeleton_bindingAPI.GetInheritedAnimationSource())
        #skelroot_bindingAPI = UsdSkel.BindingAPI.Apply(skelroot_prim)
        #skelroot_bindingAPI.GetAnimationSourceRel().SetTargets([animation_flat_path])
        #await wait_for_update()
        #self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_flat_prim)
        #timeline.set_current_time(0.0)
        #await wait_for_update()
        #await self.capture_and_compare(self.TEST_PATH, "01_skelanim_4_animation_binding_on_root_0.png")
        #timeline.set_current_time(0.5)
        #await wait_for_update()
        #await self.capture_and_compare(self.TEST_PATH, "01_skelanim_4_animation_binding_on_root_1.png")

        #skelroot_bindingAPI.GetAnimationSourceRel().ClearTargets(False)
        #self.assertTrue(not skelroot_bindingAPI.GetInheritedAnimationSource())

        ##Test animation switch to outside root animation in session layer
        #with Usd.EditContext(stage, session_layer):
        #    skelroot_bindingAPI.GetAnimationSourceRel().SetTargets([animation_flat_path])
        #    await wait_for_update()
        #    self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_flat_prim)

        #    await wait_for_update()
        #    await self.capture_and_compare(self.TEST_PATH, "01_skelanim_4_animation_binding_on_root_sessionlayer_0.png")
        #    timeline.set_current_time(0.5)
        #    await wait_for_update()
        #    await self.capture_and_compare(self.TEST_PATH, "01_skelanim_4_animation_binding_on_root_sessionlayer_1.png")

        #    stage.RemovePrim(skeleton_path)

        timeline.set_auto_update(True)
        timeline.stop()


    async def test_02_skel_anim_update(self):
        """
        Test hydra skel - skel animation update
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skeleton_path = "/Root/group1/joint1"
        skeleton_prim = stage.GetPrimAtPath(skeleton_path)
        skeleton_bindingAPI = UsdSkel.BindingAPI(skeleton_prim)
        skeleton = UsdSkel.Skeleton(skeleton_prim)
        animation_static_path = "/Root/group1/joint1/Animation_Static"
        animation_static_prim = stage.GetPrimAtPath(animation_static_path)
        animation_static = UsdSkel.Animation(animation_static_prim)
        animation_path = "/Root/group1/joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        animation = UsdSkel.Animation(animation_prim)
        animation_flat_path = "/Root/group1/joint1/Animation_Flat"
        animation_flat_prim = stage.GetPrimAtPath(animation_flat_path)
        aniamtion_flat = UsdSkel.Animation(animation_flat_prim)
        animation_outside_path = "/ZAnimation"
        animation_outside_prim = stage.GetPrimAtPath(animation_outside_path)
        aniamtion_outside = UsdSkel.Animation(animation_outside_prim)

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()

        #Test SRT/bsweights update
        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_static_path])
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_static_prim)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_0_animation_update_0.png")
        timecode = Usd.TimeCode(0.5 * stage.GetTimeCodesPerSecond())
        with Sdf.ChangeBlock():
            _update_animation(animation, animation_static, timecode)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_0_animation_update_1.png") #non play test
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(1.0)
        await omni.kit.app.get_app().next_update_async()
        timecode = Usd.TimeCode(timeline.get_current_time() * stage.GetTimeCodesPerSecond())
        with Sdf.ChangeBlock():
            _update_animation(animation, animation_static, timecode)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_0_animation_update_2.png") # play test
        timeline.set_auto_update(True)
        timeline.stop()

        timeline.set_current_time(0.0)
        await omni.kit.app.get_app().next_update_async()
        with Sdf.ChangeBlock():
            _update_animation(animation, animation_static, Usd.TimeCode.Default())
        #Test SRT/bsweights update in session layer
        with Usd.EditContext(stage, session_layer):
            self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_static_prim)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_0_animation_update_sessionlayer_0.png")
            timecode = Usd.TimeCode(0.5 * stage.GetTimeCodesPerSecond())
            with Sdf.ChangeBlock():
                _update_animation(animation, animation_static, timecode)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_0_animation_update_sessionlayer_1.png") #non play test
            stage.RemovePrim(animation_static_path)

            timeline.set_current_time(1.0)
            await omni.kit.app.get_app().next_update_async()
            timecode = Usd.TimeCode(timeline.get_current_time() * stage.GetTimeCodesPerSecond())
            with Sdf.ChangeBlock():
                _update_animation(animation, animation_static, timecode)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_2_animation_update_sessionlayer_2.png") # play test
            timeline.set_auto_update(True)
            timeline.stop()
            await omni.kit.app.get_app().next_update_async()

            stage.RemovePrim(animation_static_path)

        ###Test joint change
        ##TODO 03: omnihydra does not support joints/blendshapes change in skelanimation prim
        ##TODO 04: omnihydra does not support blendshapes (target name) change in skelanimation prim (not in test)
        #with Sdf.ChangeBlock():
            #_update_animation(animation, animation_static, Usd.TimeCode.Default())
        #await wait_for_update()
        #await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_3_joint_update_0.png")
        #timecode = Usd.TimeCode(1.0 * stage.GetTimeCodesPerSecond())
        #trans = animation.GetTranslationsAttr().Get(timecode)
        #rots = animation.GetRotationsAttr().Get(timecode)
        #scales = animation.GetScalesAttr().Get(timecode)
        #joints = animation.GetJointsAttr().Get()
        #new_joints = joints.__getitem__(slice(0,3,1))
        #new_trans = trans.__getitem__(slice(0,3,1))
        #new_rots = rots.__getitem__(slice(0,3,1))
        #new_scales = scales.__getitem__(slice(0,3,1))
        #animation_static.GetJointsAttr().Set(new_joints)
        #animation_static.GetTranslationsAttr().Set(new_trans)
        #animation_static.GetRotationsAttr().Set(new_rots)
        #animation_static.GetScalesAttr().Set(new_scales)
        #await wait_for_update()
        #await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_3_joint_update_1.png")
        #new_joints = joints.__getitem__(slice(0,2,1))
        #new_trans = trans.__getitem__(slice(0,2,1))
        #new_rots = rots.__getitem__(slice(0,2,1))
        #new_scales = scales.__getitem__(slice(0,2,1))
        #animation_static.GetJointsAttr().Set(new_joints)
        #animation_static.GetTranslationsAttr().Set(new_trans)
        #animation_static.GetRotationsAttr().Set(new_rots)
        #animation_static.GetScalesAttr().Set(new_scales)
        #await wait_for_update()
        #await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_3_joint_update_2.png")
        #animation_static.GetJointsAttr().Set(joints)
        #animation_static.GetTranslationsAttr().Set(trans)
        #animation_static.GetRotationsAttr().Set(rots)
        #animation_static.GetScalesAttr().Set(scales)
        #await wait_for_update()

        ###Test joint change in session layer
        #animation_static.GetJointsAttr().Set(joints)
        #with Sdf.ChangeBlock():
            #_update_animation(animation, animation_static, Usd.TimeCode.Default())
        #await wait_for_update()
        #await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_3_joint_update_sessionlayer_0.png")
        #with Usd.EditContext(stage, session_layer):
        #    timecode = Usd.TimeCode(1.0 * stage.GetTimeCodesPerSecond())
        #    trans = animation.GetTranslationsAttr().Get(timecode)
        #    rots = animation.GetRotationsAttr().Get(timecode)
        #    scales = animation.GetScalesAttr().Get(timecode)
        #    joints = animation.GetJointsAttr().Get()
        #    new_joints = joints.__getitem__(slice(0,3,1))
        #    new_trans = trans.__getitem__(slice(0,3,1))
        #    new_rots = rots.__getitem__(slice(0,3,1))
        #    new_scales = scales.__getitem__(slice(0,3,1))
        #    animation_static.GetJointsAttr().Set(new_joints)
        #    animation_static.GetTranslationsAttr().Set(new_trans)
        #    animation_static.GetRotationsAttr().Set(new_rots)
        #    animation_static.GetScalesAttr().Set(new_scales)
        #    await wait_for_update()
        #    await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_3_joint_update_sessionlayer_1.png")
        #    new_joints = joints.__getitem__(slice(0,2,1))
        #    new_trans = trans.__getitem__(slice(0,2,1))
        #    new_rots = rots.__getitem__(slice(0,2,1))
        #    new_scales = scales.__getitem__(slice(0,2,1))
        #    animation_static.GetJointsAttr().Set(new_joints)
        #    animation_static.GetTranslationsAttr().Set(new_trans)
        #    animation_static.GetRotationsAttr().Set(new_rots)
        #    animation_static.GetScalesAttr().Set(new_scales)
        #    await wait_for_update()
        #    await self.capture_and_compare(self.TEST_PATH, "02_skelanim_update_3_joint_update_sessionlayer_2.png")
        #    animation_static.GetJointsAttr().Set(joints)
        #    animation_static.GetTranslationsAttr().Set(trans)
        #    animation_static.GetRotationsAttr().Set(rots)
        #    animation_static.GetScalesAttr().Set(scales)
        #    await wait_for_update()

        #    stage.RemovePrim(animation_static_path)


    async def test_03_skel_anim_create_delete(self):
        """
        Test hydra skel - skel animation update
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skeleton_path = "/Root/group1/joint1"
        skeleton_prim = stage.GetPrimAtPath(skeleton_path)
        skeleton_bindingAPI = UsdSkel.BindingAPI(skeleton_prim)
        skeleton = UsdSkel.Skeleton(skeleton_prim)
        animation_static_path = "/Root/group1/joint1/Animation_Static"
        animation_static_prim = stage.GetPrimAtPath(animation_static_path)
        animation_static = UsdSkel.Animation(animation_static_prim)
        animation_path = "/Root/group1/joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        animation = UsdSkel.Animation(animation_prim)
        animation_flat_path = "/Root/group1/joint1/Animation_Flat"
        animation_flat_prim = stage.GetPrimAtPath(animation_flat_path)
        aniamtion_flat = UsdSkel.Animation(animation_flat_prim)
        animation_outside_path = "/ZAnimation"
        animation_outside_prim = stage.GetPrimAtPath(animation_outside_path)
        aniamtion_outside = UsdSkel.Animation(animation_outside_prim)

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_current_time(0.0)
        await omni.kit.app.get_app().next_update_async()

        timecode = Usd.TimeCode(1.0 * stage.GetTimeCodesPerSecond())
        trans = animation.GetTranslationsAttr().Get(timecode)
        rots = animation.GetRotationsAttr().Get(timecode)
        scales = animation.GetScalesAttr().Get(timecode)
        joints = animation.GetJointsAttr().Get()
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "03_skelanim_create_delete_0_animation_create_0.png")

        animation_new_path = "/ZAnimation_New"
        animation_new = UsdSkel.Animation.Define(stage, animation_new_path)
        animation_new_prim = animation_new.GetPrim()
        await wait_for_update()
        self.assertTrue(animation_new_prim)
        with Sdf.ChangeBlock():
            animation_new.GetJointsAttr().Set(joints)
            animation_new.GetTranslationsAttr().Set(trans)
            animation_new.GetRotationsAttr().Set(rots)
            animation_new.GetScalesAttr().Set(scales)
        skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_new_path])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "03_skelanim_create_delete_0_animation_create_1.png")

        stage.RemovePrim(animation_new_path)
        skeleton_bindingAPI.GetAnimationSourceRel().ClearTargets(False) #This line should supposed not to be needed #TODO 05 when skelanimation become invalid, skinning result should retrive restTransform
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "03_skelanim_create_delete_1_animation_delete_0.png")

        with Usd.EditContext(stage, session_layer):
            animation_new = UsdSkel.Animation.Define(stage, animation_new_path)
            animation_new_prim = animation_new.GetPrim()
            await wait_for_update()
            self.assertTrue(animation_new_prim)
            with Sdf.ChangeBlock():
                animation_new.GetJointsAttr().Set(joints)
                animation_new.GetTranslationsAttr().Set(trans)
                animation_new.GetRotationsAttr().Set(rots)
                animation_new.GetScalesAttr().Set(scales)
            skeleton_bindingAPI.GetAnimationSourceRel().SetTargets([animation_new_path])
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "03_skelanim_create_delete_2_animation_create_sessionlayer_0.png")
            with Sdf.ChangeBlock():
                stage.RemovePrim(animation_new_path)
                skeleton_bindingAPI.GetAnimationSourceRel().ClearTargets(False) #This line should supposed not to be needed #TODO 05 when skelanimation become invalid, skinning result should retrive restTransform
                skel_root_prim = stage.GetPrimAtPath("/Root/group1")
                skel_root = UsdSkel.Root(skel_root_prim)
                visible = skel_root.GetVisibilityAttr().Set("invisible")
            await wait_for_update() # test repopulate crash due to resync by update visibility in skelroot.
            skel_root.GetVisibilityAttr().Set("inherited")
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "03_skelanim_create_delete_3_animation_delete_with_root_recync_sessionlayer_0.png")

    #This test is temporarily since we cannot dynamically update animation source on skelroot's bindAPI in omnihydra yet.
    async def test_04_skel_anim_on_skel_root(self):
        """
        Test hydra skel - skel anim on root
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder_root.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skelroot_path = "/Root/group1"
        skelroot_prim = stage.GetPrimAtPath(skelroot_path)
        skelroot = UsdSkel.Root(skelroot_prim)
        skelroot_bindingAPI = UsdSkel.BindingAPI(skelroot_prim)
        skeleton_path = "/Root/group1/joint1"
        skeleton_prim = stage.GetPrimAtPath(skeleton_path)
        skeleton_bindingAPI = UsdSkel.BindingAPI(skeleton_prim)
        skeleton = UsdSkel.Skeleton(skeleton_prim)
        animation_static_path = "/Root/group1/joint1/Animation_Static"
        animation_static_prim = stage.GetPrimAtPath(animation_static_path)
        animation_static = UsdSkel.Animation(animation_static_prim)
        animation_path = "/Root/group1/joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        animation = UsdSkel.Animation(animation_prim)
        animation_flat_path = "/Root/group1/joint1/Animation_Flat"
        animation_flat_prim = stage.GetPrimAtPath(animation_flat_path)
        aniamtion_flat = UsdSkel.Animation(animation_flat_prim)
        animation_outside_path = "/ZAnimation"
        animation_outside_prim = stage.GetPrimAtPath(animation_outside_path)
        aniamtion_outside = UsdSkel.Animation(animation_outside_prim)

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(0.0)
        await omni.kit.app.get_app().next_update_async()

        # pure skeleton test
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "04_skelanim_on_root_0_skeleton_0.png")

        #Test animation
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
        timeline.set_current_time(0.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "04_skelanim_on_root_1_animation_0.png")
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "04_skelanim_on_root_1_animation_1.png")
        timeline.set_current_time(1.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "04_skelanim_on_root_1_animation_2.png")

        timeline.set_auto_update(True)
        timeline.stop()
        await omni.kit.app.get_app().next_update_async()

    async def test_05_skel_anim_update_restTransforms(self):
        """
        Test hydra skel - update restTransforms
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder.usda")
        # self.set_settings(postLoadSkelTestSettings)
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skelroot_path = "/Root/group1"
        skelroot_prim = stage.GetPrimAtPath(skelroot_path)
        skelroot = UsdSkel.Root(skelroot_prim)
        skelroot_bindingAPI = UsdSkel.BindingAPI(skelroot_prim)
        skeleton_path = "/Root/group1/joint1"
        skeleton_prim = stage.GetPrimAtPath(skeleton_path)
        skeleton_bindingAPI = UsdSkel.BindingAPI(skeleton_prim)
        skeleton = UsdSkel.Skeleton(skeleton_prim)
        animation_static_path = "/Root/group1/joint1/Animation_Static"
        animation_static_prim = stage.GetPrimAtPath(animation_static_path)
        animation_static = UsdSkel.Animation(animation_static_prim)
        animation_path = "/Root/group1/joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        animation = UsdSkel.Animation(animation_prim)
        animation_flat_path = "/Root/group1/joint1/Animation_Flat"
        animation_flat_prim = stage.GetPrimAtPath(animation_flat_path)
        aniamtion_flat = UsdSkel.Animation(animation_flat_prim)
        animation_outside_path = "/ZAnimation"
        animation_outside_prim = stage.GetPrimAtPath(animation_outside_path)
        aniamtion_outside = UsdSkel.Animation(animation_outside_prim)

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(0.0)
        await omni.kit.app.get_app().next_update_async()

        # pure skeleton test
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_0_skeleton_0.png")

        #Test animation
        self.assertTrue(not skeleton_bindingAPI.GetInheritedAnimationSource())
        timecode = Usd.TimeCode.Default()
        poses = animation.GetTransforms(timecode)
        skeleton.GetRestTransformsAttr().Set(poses)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_1_animation_0.png")
        timecode = Usd.TimeCode(0.5 * stage.GetTimeCodesPerSecond())
        poses = animation.GetTransforms(timecode)
        skeleton.GetRestTransformsAttr().Set(poses)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_1_animation_1.png")
        timecode = Usd.TimeCode(1.0 * stage.GetTimeCodesPerSecond())
        poses = animation.GetTransforms(timecode)
        skeleton.GetRestTransformsAttr().Set(poses)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_1_animation_2.png")

        with Usd.EditContext(stage, session_layer):
            timecode = Usd.TimeCode.Default()
            poses = animation.GetTransforms(timecode)
            skeleton.GetRestTransformsAttr().Set(poses)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_1_animation_session_layer_0.png")
            timecode = Usd.TimeCode(0.5 * stage.GetTimeCodesPerSecond())
            poses = animation.GetTransforms(timecode)
            skeleton.GetRestTransformsAttr().Set(poses)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_1_animation_session_layer_1.png")
            timecode = Usd.TimeCode(1.0 * stage.GetTimeCodesPerSecond())
            poses = animation.GetTransforms(timecode)
            skeleton.GetRestTransformsAttr().Set(poses)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "05_skelanim_update_restTransforms_1_animation_session_layer_2.png")

        timeline.set_auto_update(True)
        timeline.stop()
        await omni.kit.app.get_app().next_update_async()


    async def test_06_skel_anim_reference(self):
        """
        Test hydra skel - skel anim reference
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder_ref.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skelroot_path = "/Root/group1"
        skelroot_prim = stage.GetPrimAtPath(skelroot_path)
        skelroot = UsdSkel.Root(skelroot_prim)
        skelroot_bindingAPI = UsdSkel.BindingAPI(skelroot_prim)
        skeleton_path = "/Root/group1/joint1"
        skeleton_prim = stage.GetPrimAtPath(skeleton_path)
        skeleton_bindingAPI = UsdSkel.BindingAPI(skeleton_prim)
        skeleton = UsdSkel.Skeleton(skeleton_prim)
        animation_static_path = "/Root/group1/joint1/Animation_Static"
        animation_static_prim = stage.GetPrimAtPath(animation_static_path)
        animation_static = UsdSkel.Animation(animation_static_prim)
        animation_path = "/Root/group1/joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        animation = UsdSkel.Animation(animation_prim)
        animation_flat_path = "/Root/group1/joint1/Animation_Flat"
        animation_flat_prim = stage.GetPrimAtPath(animation_flat_path)
        aniamtion_flat = UsdSkel.Animation(animation_flat_prim)
        animation_outside_path = "/ZAnimation"
        animation_outside_prim = stage.GetPrimAtPath(animation_outside_path)
        aniamtion_outside = UsdSkel.Animation(animation_outside_prim)

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(0.0)
        await omni.kit.app.get_app().next_update_async()

        # pure skeleton tests
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "06_skelanim_reference_0_skeleton_0.png")

        # Test animation
        self.assertTrue(skeleton_bindingAPI.GetInheritedAnimationSource() == animation_prim)
        refs = animation_prim.GetReferences()
        refs.SetReferences([Sdf.Reference(assetPath="./assets/skelcylinder_anim_flat.usda")])
        timeline.set_current_time(0.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "06_skelanim_reference_1_animation_reference_0.png")
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "06_skelanim_reference_1_animation_reference_1.png")
        timeline.set_current_time(1.0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "06_skelanim_reference_1_animation_reference_2.png")

        timeline.set_auto_update(True)
        timeline.stop()


    async def test_07_skel_mesh_material_switch(self):
        """
        Test hydra skel - skel anim reference
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder_material_test.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skelmesh1_path = "/Root/BLENDWEIGHT_group1/BLENDWEIGHT_pCylinder1"
        skelmesh1_prim = stage.GetPrimAtPath(skelmesh1_path)
        skelmesh2_path = "/Root/BLENDWEIGHT_group1/BLENDWEIGHT_pCylinder2"
        skelmesh2_prim = stage.GetPrimAtPath(skelmesh2_path)
        print(skelmesh2_prim)
        material_red_path = "/Root/Looks/PreviewSurface_Red"
        material_blue_path = "/Root/Looks/PreviewSurface_Blue"

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(0.0)
        await wait_for_update()

        # pure skeleton tests
        timeline.set_current_time(0.1)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "07_skel_mesh_material_switch_0.png")

        # Test animations
        omni.kit.commands.execute(
            "BindMaterial",
            prim_path=Sdf.Path(skelmesh1_path),
            material_path=Sdf.Path(material_blue_path),
            strength=UsdShade.Tokens.weakerThanDescendants,
        )
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "07_skel_mesh_material_switch_1.png")

        omni.kit.commands.execute(
            "BindMaterial",
            prim_path=Sdf.Path(skelmesh2_path),
            material_path=Sdf.Path(material_red_path),
            strength=UsdShade.Tokens.weakerThanDescendants,
        )
        timeline.set_current_time(0.6)
        await wait_for_update()
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "07_skel_mesh_material_switch_2.png")

        imageable = UsdGeom.Imageable(skelmesh2_prim)
        imageable.MakeInvisible()
        timeline.set_current_time(0.7)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "07_skel_mesh_material_switch_3.png")

        timeline.set_auto_update(True)
        timeline.stop()

    async def test_08_skel_mesh_animation_rename(self):
        """
        Test hydra skel - skel anim reference
        """
        self.set_settings(testSettings)
        super().open_usd("hydra/skel/skelcylinder_material_test.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadSkelTestSettings)

        stage = self.ctx.get_stage()
        skelmesh1_path = "/Root/BLENDWEIGHT_group1/BLENDWEIGHT_pCylinder1"
        skelmesh1_prim = stage.GetPrimAtPath(skelmesh1_path)
        skelmesh2_path = "/Root/BLENDWEIGHT_group1/BLENDWEIGHT_pCylinder2"
        skelmesh2_prim = stage.GetPrimAtPath(skelmesh2_path)
        animation_path = "/Root/BLENDWEIGHT_group1/BLENDWEIGHT_joint1/Animation"
        animation_prim = stage.GetPrimAtPath(animation_path)
        new_animation_path = "/Root/BLENDWEIGHT_group1/BLENDWEIGHT_joint1/Animation2"

        session_layer = stage.GetSessionLayer()
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        timeline.set_auto_update(False)
        timeline.set_current_time(0.0)
        await wait_for_update()

        # pure skeleton tests
        timeline.set_current_time(0.5)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "08_skel_mesh_animation_rename_0.png")
        old_prim_name = Sdf.Path(animation_path)
        move_dict = {old_prim_name: new_animation_path}
        omni.kit.commands.execute("MovePrims", paths_to_move=move_dict, destructive=False)

        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "08_skel_mesh_animation_rename_1.png")

        omni.kit.undo.undo()

        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "08_skel_mesh_animation_rename_2.png")

        # Test animations
        omni.kit.commands.execute(
            "DeletePrims",
            paths=[Sdf.Path(animation_path)]
        )

        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "08_skel_mesh_animation_rename_3.png")

        omni.kit.undo.undo()
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "08_skel_mesh_animation_rename_4.png")

        timeline.set_auto_update(True)
        timeline.stop()
