## 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.test
from .test_common import RtxTest, testSettings, postLoadTestSettings, wait_for_update
from pxr import UsdGeom, UsdLux, Gf, Sdf

class TestRtxDomelight(RtxTest):
    TEST_PATH = "domelight"
    DOMELIGHT_PRIM_PATH = "/World/DomeLight"

    def create_mesh(self):
        box = UsdGeom.Mesh.Define(self.ctx.get_stage(), "/World/box")
        box.CreatePointsAttr([(-50, -50, -50), (50, -50, -50), (-50, -50, 50), (50, -50, 50),
                              (-50, 50, -50), (50, 50, -50), (50, 50, 50), (-50, 50, 50)])
        box.CreateFaceVertexCountsAttr([4, 4, 4, 4, 4, 4])
        box.CreateFaceVertexIndicesAttr([0, 1, 3, 2, 0, 4, 5, 1, 1, 5, 6, 3, 2, 3, 6, 7, 0, 2, 7, 4, 4, 7, 6, 5])
        box.CreateSubdivisionSchemeAttr("none")
        return box

    def create_dome_light(self, name=DOMELIGHT_PRIM_PATH):
        omni.kit.commands.execute(
            "CreatePrim",
            prim_path=name,
            prim_type="DomeLight",
            select_new_prim=False,
            # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
            attributes={
                UsdLux.Tokens.inputsIntensity: 1,
                UsdLux.Tokens.inputsTextureFormat: UsdLux.Tokens.latlong,
                UsdLux.Tokens.inputsTextureFile: "daytime.hdr",
                UsdGeom.Tokens.visibility: UsdGeom.Tokens.inherited,
            } if hasattr(UsdLux.Tokens, 'inputsIntensity') else
            {
                UsdLux.Tokens.intensity: 1,
                UsdLux.Tokens.textureFormat: UsdLux.Tokens.latlong,
                UsdLux.Tokens.textureFile: "daytime.hdr",
                UsdGeom.Tokens.visibility: UsdGeom.Tokens.inherited,
            },
            create_default_xform=True,
        )
        dome_light_prim = self.ctx.get_stage().GetPrimAtPath(name)
        return dome_light_prim

    async def setUp(self):
        await self.setUp_internal()
        print("RTX DomeLight Tests Setup")

        self.set_settings(testSettings)
        super().open_usd("hydra/dome_materials.usda")
        await omni.kit.app.get_app().next_update_async()

        self.set_settings(postLoadTestSettings)
        my_settings = {
            "/rtx/pathtracing/lightcache/cached/enabled": False,
            "/rtx/raytracing/lightcache/spatialCache/enabled" : False,
            "/rtx-transient/resourcemanager/genMipsForNormalMaps" : False,
            "/rtx-transient/resourcemanager/texturestreaming/async" : False,
            "/rtx-transient/samplerFeedbackTileSize" : 1,
            "/rtx/post/aa/op" : 0, # 0 = None, 2 = FXAA
            "/rtx/directLighting/sampledLighting/enabled" : False,
            "/rtx/reflections/sampledLighting/enabled" : False, # LTC gives consistent lighting
            # perMaterialSyncLoads: Very important, otherwise the material updates for the domelight
            # materials (MDLs) are not sync-ed and updated properly.
            "/rtx/hydra/perMaterialSyncLoads" : True,
        }
        self.set_settings(my_settings)

    async def test_domelight_material_assignment(self):
        """
        Test Domelight material assignment
        """
        looksPath = "/World/Looks/"

        # Scene Setup
        self.set_settings({"/rtx/domeLight/baking/resolution": "1024"})
        # We could show some minimal mesh but that might create false positives in the lighting.
        #self.create_mesh()
        dome_prim = self.create_dome_light()

        # The loaded domelight had a texture assigned check it:
        # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
        if hasattr(UsdLux.Tokens, 'inputsIntensity'):
            dome_prim.GetAttribute(UsdLux.Tokens.inputsIntensity).Set(100)
        else:
            dome_prim.GetAttribute(UsdLux.Tokens.intensity).Set(100)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "01_texture_assigned.png")

        # Bind a material to the domelight
        # This material shows the emission direction in a color-coded way.
        omni.kit.commands.execute('BindMaterial',
            material_path=looksPath + "dome_emission_direction",
            prim_path=[self.DOMELIGHT_PRIM_PATH],
            strength=['weakerThanDescendants'])
        # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
        if hasattr(UsdLux.Tokens, 'inputsIntensity'):
            dome_prim.GetAttribute(UsdLux.Tokens.inputsIntensity).Set(1)
        else:
            dome_prim.GetAttribute(UsdLux.Tokens.intensity).Set(1)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "02_emission_direction_mat.png")

        # Bind a different material to the domelight
        omni.kit.commands.execute('BindMaterial',
            material_path='/World/Looks/dome_gridspherejulia',
            prim_path=['/World/DomeLight'],
            strength=['weakerThanDescendants'])
        # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
        if hasattr(UsdLux.Tokens, 'inputsIntensity'):
            dome_prim.GetAttribute(UsdLux.Tokens.inputsIntensity).Set(1)
        else:
            dome_prim.GetAttribute(UsdLux.Tokens.intensity).Set(1)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "03_gridspherejulia_mat.png")

        # Set a different baking resolution
        self.set_settings({"/rtx/domeLight/baking/resolution": "256"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "04_baking_resolution.png")

        # Testing domelight per pixel evaluation:
        # JIRA OM-49492
        self.set_settings({
            "/rtx/pathtracing/domeLight/primaryRaysEvaluateDomelightMdlDirectly" : True,
            "/rtx/rendermode" : 'PathTracing',
            "/rtx/pathtracing/spp" : 1,
            "/rtx/pathtracing/totalSpp" : 1,
        })
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "04a_mdl_direct_evaluation.png")

        # Create another domelight (JIRA OM-19501)
        # The visible domelight can be any domelight if multiple domelights are in the scene.
        # Internally it depends which one is in the domelight buffer[0].
        # Therefore, we need to explicitly disable the first one to see the result of the second.
        self.set_settings({
            "/rtx/rendermode" : 'RaytracedLighting',
        })
        dome2_path = "/World/DomeLight_2"
        dome_prim2 = self.create_dome_light(dome2_path)
        # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
        if hasattr(UsdLux.Tokens, 'inputsIntensity'):
            dome_prim2.GetAttribute(UsdLux.Tokens.inputsIntensity).Set(100)
        else:
            dome_prim2.GetAttribute(UsdLux.Tokens.intensity).Set(100)
        omni.kit.commands.execute('ChangeProperty',
            prop_path=Sdf.Path('/World/DomeLight_2.xformOp:rotateXYZ'),
            value=Gf.Vec3d(270.0, -90.0, 0.0),
            prev=Gf.Vec3d(270.0, 0.0, 0.0))

        # Disable the first dome light
        dome_prim.GetAttribute(UsdGeom.Tokens.visibility).Set(UsdGeom.Tokens.invisible)

        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "05_second_domelight.png")
