## 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.
##
import omni.kit.app
import omni.kit.commands
import omni.kit.undo
import omni.kit.test
import omni.timeline
from .test_common import RtxTest, testSettings, postLoadTestSettings, set_transform_helper, wait_for_update
from pxr import Sdf, Gf, UsdGeom, UsdShade


class RtxHydraTest(RtxTest):
    """
    Base class for hydra tests
    """
    # override resolution
    WINDOW_SIZE = (512, 512)
    TEST_PATH = "hydra"
    PRIM_PATH = ""

    async def setUp(self):
        await self.setUp_internal()
        self.set_settings(testSettings)
        omni.usd.get_context().new_stage()
        self.add_dir_light()
        await omni.kit.app.get_app().next_update_async()  # Wait stage loading
        self.set_settings(postLoadTestSettings)

    def create_geometry(self, name=PRIM_PATH):
        pass

    def create_geometry_with_floor(self):
        self.add_floor()
        self.create_geometry()
        set_transform_helper(self.PRIM_PATH, translate=Gf.Vec3d(0, 100, 0))

    async def geometry_with_attribute(self, attrName, attrType, attrVal, shouldWait=False):
        geom = self.create_geometry()
        if shouldWait:
            await wait_for_update()
        attr = geom.GetPrim().CreateAttribute(attrName, attrType)
        self.assertTrue(attr)
        attr.Set(attrVal)
        await wait_for_update()
        return geom

    async def set_attribute_timesampled(self, geom, attrName, attrType, attrVals, goldens):
        attr = geom.GetPrim().CreateAttribute(attrName, attrType)
        self.assertTrue(attr)
        for i, v in enumerate(attrVals):
            attr.Set(v, i)
        self.assertTrue(len(attrVals) == len(goldens))
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_target_framerate(timeline.get_time_codes_per_seconds())
        timeline.play()
        timeline.set_auto_update(False)
        for i, g in enumerate(goldens):
            timeline.set_current_time(i)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, g)
        timeline.set_auto_update(True)
        timeline.stop()

    async def attribute_test_case(self, createGeom, attrName, attrType, attrVals, goldens):
        self.assertTrue(len(attrVals) > 1)
        self.assertTrue(len(attrVals) == len(goldens))
        geom = createGeom()
        attr = geom.GetPrim().CreateAttribute(attrName, attrType)
        self.assertTrue(attr)
        # Test updating attribute
        for i, g in enumerate(goldens):
            attr.Set(attrVals[i])
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, g)

        # Test timesampled attribute
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        geom = createGeom()
        await self.set_attribute_timesampled(geom, attrName, attrType, attrVals, goldens)

    async def toggle_primvar(self, geomPath, primvarName, primvarType=Sdf.ValueTypeNames.Bool, primvarVal=True,
                             baseGoldenName="boxAndFloor.png", threshould=1e-5):
        prim = self.ctx.get_stage().GetPrimAtPath(geomPath)
        attr = UsdGeom.PrimvarsAPI(prim).CreatePrimvar(primvarName, primvarType)
        self.assertTrue(attr)
        await wait_for_update()
        self.assertFalse(attr.Get() == primvarVal)
        await self.capture_and_compare(self.TEST_PATH, baseGoldenName)
        attr.Set(primvarVal)
        await wait_for_update()
        self.assertTrue(attr.Get() == primvarVal)
        await self.capture_and_compare(self.TEST_PATH, primvarName + ".png", threshould)

    async def visibility(self):
        prim = self.create_geometry()
        attr = prim.GetVisibilityAttr()
        self.assertTrue(attr)
        
        for i in range(2):
            attr.Set(UsdGeom.Tokens.inherited)
            await wait_for_update()
            self.assertTrue(attr.Get() == UsdGeom.Tokens.inherited)
            await self.capture_and_compare(self.TEST_PATH, "geom.png")
            
            attr.Set(UsdGeom.Tokens.invisible)
            await wait_for_update()
            self.assertTrue(attr.Get() == UsdGeom.Tokens.invisible)
            await self.capture_and_compare(self.TEST_PATH, "empty.png")

    async def cull_style(self):
        # Not yet supported by hydra
        pass

    async def double_sided(self):
        pass

    async def single_sided(self):
        self.set_settings({"/rtx/debugView/target": "normal", 
                           "/rtx/hydra/faceCulling/enabled": True})
        geom = self.create_geometry()
        attr = geom.GetPrim().CreateAttribute("singleSided", Sdf.ValueTypeNames.Bool)
        self.assertTrue(attr)
        self.set_camera(cameraPos=(0, 0, 0), targetPos=(-500, -500, -500))
        await wait_for_update()
        self.assertFalse(attr.Get())
        await self.capture_and_compare(self.TEST_PATH, "doubleSided.png")
        attr.Set(True)
        await wait_for_update()
        self.assertTrue(attr.Get())
        await self.capture_and_compare(self.TEST_PATH, "singleSided.png")

    async def matte(self):
        self.set_settings({"/rtx/matteObject/enabled": True,
                           "/rtx/post/backgroundZeroAlpha/enabled": True,
                           "/rtx/post/backgroundZeroAlpha/backgroundComposite": True,
                           "/rtx/post/backgroundZeroAlpha/backgroundDefaultColor": (0, 0, 0)})
        await self.toggle_primvar("/World/Floor", "isMatteObject")

    async def is_picked(self):
        # This is not worked when kit is not focused. Skip tests while solution wouldn't be found
        self.skipTest("Test is not worked correctly yet!")
        #viewport = omni.kit.viewport_legacy.acquire_viewport_interface().get_viewport_window(None)
        #isPicked = False 
        # for i in range(10):
        #     set_cursor_position((int(self.WINDOW_SIZE[0] / 2), int(self.WINDOW_SIZE[1] / 2)))
        #     await omni.kit.app.get_app().next_update_async()
        #     ret = viewport.get_hovered_world_position()
        #     isPicked |= ret[0]
        #return isPicked

    async def pickable(self):
        geom = self.create_geometry()
        self.assertTrue(await self.is_picked())
        self.ctx.set_pickable(self.PRIM_PATH, False)
        # set_pickable is not catched correctly by hydra yet. So we need to update any primvar to force update instance
        attr = geom.GetPrim().CreateAttribute("singleSided", Sdf.ValueTypeNames.Bool)
        attr.Set(True)
        await wait_for_update()
        self.assertFalse(await self.is_picked())

    async def transform_translate(self):
        self.create_geometry()
        set_transform_helper(self.PRIM_PATH, translate=Gf.Vec3d(0, 0, 150))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "translateZ.png")

    async def transform_rotate(self):
        self.create_geometry()
        set_transform_helper(self.PRIM_PATH, euler=Gf.Vec3d(0, 45, 0))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "rotateY.png")

    async def transform_scale(self):
        self.create_geometry()
        set_transform_helper(self.PRIM_PATH, scale=Gf.Vec3d(2, 1, 1))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "scaleX.png")

    async def transform_combined(self, shouldWait=False):
        self.create_geometry()
        if shouldWait:
            await wait_for_update()
        set_transform_helper(self.PRIM_PATH, translate=Gf.Vec3d(0, 0, 150), euler=Gf.Vec3d(0, 45, 0),
                             scale=Gf.Vec3d(2, 1, 1))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "tranlateZrotateYscaleX.png")

    async def transform_all(self):
        await self.transform_translate()
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        await self.transform_rotate()
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        await self.transform_scale()
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)

        await self.transform_combined(shouldWait=True)
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        await wait_for_update()
        await self.transform_combined(shouldWait=False)


class RtxHydraMaterialsTest(RtxHydraTest):
    TEST_PATH = ""
    PRIM_PATH = ""

    async def setUp(self):
        await self.setUp_internal()
        self.set_settings(testSettings)
        super().open_usd("hydra/UsdShade.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

    def create_geometry(self, name):
        pass

    def bind_material(self, geom, material_path):
        binding_api = UsdShade.MaterialBindingAPI(geom)
        material_prim = self.ctx.get_stage().GetPrimAtPath(material_path)
        material = UsdShade.Material(material_prim)
        binding_api.Bind(material, UsdShade.Tokens.weakerThanDescendants)


class RtxHydraInstancerTest(RtxHydraTest):
    # TODO full cover scene instancer with tests
    TEST_PATH = ""
    GEOM_PATH = ""
    PRIM_PATH = "/World/instance"
    PRIM1_PATH = "/World/instance1"
    PRIM2_PATH = "/World/instance2"

    def create_geometry(self, name):
        pass

    def create_scene_instancer(self):
        prim = UsdGeom.Xform.Define(self.ctx.get_stage(), self.PRIM_PATH).GetPrim()
        self.create_geometry(self.GEOM_PATH)
        prim1 = UsdGeom.Xform.Define(self.ctx.get_stage(), self.PRIM1_PATH).GetPrim()
        prim1.GetReferences().AddInternalReference(self.PRIM_PATH)
        prim1.SetInstanceable(True)

        omni.kit.commands.execute("CopyPrimCommand", path_from=self.PRIM1_PATH, path_to=self.PRIM2_PATH)
        self.ctx.get_selection().clear_selected_prim_paths()
        prim2 = self.ctx.get_stage().GetPrimAtPath(self.PRIM2_PATH)
        set_transform_helper(self.PRIM2_PATH, translate=Gf.Vec3d(0, 200, 0))
        return prim, prim2
    
    def remove_scene_instancer(self):
        # Teardown in reverse order of creation, to avoid temporarily leaving dangling references.
        # Alternatively, we could implement this in Sdf with a change block to batch all the removals.
        self.ctx.get_stage().RemovePrim(self.PRIM2_PATH)
        self.ctx.get_stage().RemovePrim(self.PRIM1_PATH)
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        self.ctx.get_stage().RemovePrim(self.GEOM_PATH)

    def create_point_instancer(self):
        instancer = UsdGeom.PointInstancer.Define(self.ctx.get_stage(), self.PRIM_PATH)
        instancer.CreatePositionsAttr([(-150, 0, 0), (150, 0, 0)])
        instancer.CreateProtoIndicesAttr([0, 0])
        self.create_geometry(self.GEOM_PATH)
        attr = instancer.CreatePrototypesRel()
        attr.AddTarget(self.GEOM_PATH)
        return instancer

    def remove_point_instancer(self):
        self.ctx.get_stage().RemovePrim(self.GEOM_PATH)
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)

    async def si_create(self):
        self.create_scene_instancer()
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "sceneinstancer.png", 1e-4)

    async def si_rotate(self):
        self.create_scene_instancer()
        set_transform_helper(self.GEOM_PATH, euler=Gf.Vec3d(0, 0, 45))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "si_rotate_mesh.png", 1e-4)

    async def si_all(self):
        await self.si_create()
        self.remove_scene_instancer()
        await self.si_rotate()
        self.remove_scene_instancer()

    async def pi_create(self):
        self.create_point_instancer()
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pointinstancer.png", 1e-4)

    async def pi_protoIndices(self):
        instancer = self.create_point_instancer()
        sphere = UsdGeom.Sphere.Define(self.ctx.get_stage(), self.PRIM_PATH + "/sphere")
        sphere.CreateRadiusAttr(50)
        instancer.GetPrototypesRel().AddTarget(self.PRIM_PATH + "/sphere")
        instancer.GetProtoIndicesAttr().Set([1, 0])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_2meshes.png", 1e-4)

        # Regression test for OM-32503
        instancer.GetProtoIndicesAttr().Set([1, 0, 1])
        instancer.GetPositionsAttr().Set([(-150, 0, 0), (0, 0, 0), (150, 0, 0)])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_3meshes.png", 1e-4)

    async def pi_invisibleIds(self):
        instancer = self.create_point_instancer()
        attr = instancer.CreateInvisibleIdsAttr()
        attr.Set([0])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_invisibleIds.png", 1e-4)
        attr.Set([])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pointinstancer.png", 1e-4)

    async def pi_orientations(self):
        instancer = self.create_point_instancer()
        attr = instancer.CreateOrientationsAttr()
        attr.Set([Gf.Quath(0., 0., 0.5, 1), Gf.Quath(0., 0.5, 0., 1)])
        
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_orientations.png", 1e-4)
        

    async def pi_scales(self):
        instancer = self.create_point_instancer()
        attr = instancer.CreateScalesAttr()
        attr.Set([(1., 2., 1), (1., 1., 2)])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_scales.png", 1e-4)

    # OM-30814 - verify proper refresh when authoring xforms on prototype gprims and ancestors
    async def pi_protoXform(self):
        instancer = self.create_point_instancer()
        await wait_for_update()

        # TODO: Figure out why the bug does not repro via scripting, only via TransformGizmo;
        # even scripting the entire operation within an Sdf.ChangeBlock does not trigger the
        # bug, with the fixes reverted.
        set_transform_helper(self.GEOM_PATH, scale=Gf.Vec3d(2, 1, 2))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_protoXform.png", 1e-4)
        set_transform_helper(self.GEOM_PATH, scale=Gf.Vec3d(1, 1, 1))
        await wait_for_update()

        geomPath = Sdf.Path(self.GEOM_PATH)
        xformPath = geomPath.GetParentPath().AppendChild('Xform')
        UsdGeom.Xform.Define(self.ctx.get_stage(), xformPath)
        newGeomPath = xformPath.AppendChild(geomPath.name)
        omni.kit.commands.execute("MovePrimCommand", path_from=self.GEOM_PATH, path_to=newGeomPath)
        instancer.GetPrototypesRel().ClearTargets(removeSpec=True)
        instancer.GetPrototypesRel().AddTarget(xformPath)
        await wait_for_update()
        set_transform_helper(xformPath, scale=Gf.Vec3d(1, 2, 1))
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_ancestralProtoXform.png", 1e-4)
        omni.kit.commands.execute("MovePrimCommand", path_from=newGeomPath, path_to=self.GEOM_PATH)
        instancer.GetPrototypesRel().ClearTargets(removeSpec=True)
        instancer.GetPrototypesRel().AddTarget(geomPath)
        self.ctx.get_stage().RemovePrim(xformPath)

        # OM-81682 - verify the prototype transform does not include the transform above the prototype root
        set_transform_helper(self.GEOM_PATH, scale=Gf.Vec3d(1, 1, 1))
        await wait_for_update()
        parentXformPath = geomPath.GetParentPath().AppendChild('Parent')
        UsdGeom.Xform.Define(self.ctx.get_stage(), parentXformPath)
        newXformPath = parentXformPath.AppendChild('Xform')
        UsdGeom.Xform.Define(self.ctx.get_stage(), newXformPath)
        newGeomPath = newXformPath.AppendChild(geomPath.name)
        omni.kit.commands.execute("MovePrimCommand", path_from=self.GEOM_PATH, path_to=newGeomPath)
        set_transform_helper(newXformPath, scale=Gf.Vec3d(1, 2, 1))
        set_transform_helper(parentXformPath, scale=Gf.Vec3d(2, 2, 2))
        instancer.GetPrototypesRel().ClearTargets(removeSpec=True)
        instancer.GetPrototypesRel().AddTarget(newXformPath)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pi_ancestralProtoXform.png")

        omni.kit.commands.execute("MovePrimCommand", path_from=newGeomPath, path_to=self.GEOM_PATH)
        self.ctx.get_stage().RemovePrim(parentXformPath)
        set_transform_helper(self.GEOM_PATH, scale=Gf.Vec3d(2, 1, 2))
        instancer.GetPrototypesRel().ClearTargets(removeSpec=True)
        instancer.GetPrototypesRel().AddTarget(self.GEOM_PATH)

    async def pi_all(self):
        await self.pi_create()
        self.remove_point_instancer()
        await self.pi_protoIndices()
        self.remove_point_instancer()
        await self.pi_invisibleIds()
        self.remove_point_instancer()

        # Transforms
        await self.pi_orientations()
        self.remove_point_instancer()
        await self.pi_scales()
        self.remove_point_instancer()

        self.create_point_instancer()
        await self.transform_translate()
        self.remove_point_instancer()
        self.create_point_instancer()
        await self.transform_rotate()
        self.remove_point_instancer()
        self.create_point_instancer()
        await self.transform_scale()
        self.remove_point_instancer()

        await self.pi_protoXform()
        self.remove_point_instancer()
