## Copyright (c) 2022, 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
import omni.kit.commands
import carb
import pathlib
from .test_common import RtxTest, testSettings, postLoadTestSettings, wait_for_update
from .test_common import EXTENSION_FOLDER_PATH
from pxr import Gf, Sdf, Usd

class TestRtxLightToggling(RtxTest):
    TEST_PATH = "light_toggling"

    def open_usd_scene(self, path : pathlib.Path):
        omni.usd.get_context().open_stage(str(path))

    async def setUp(self):
        await self.setUp_internal()
        carb.log_info("Setting up scene for RTX postprocessing tonemapping tests.")

        self.set_settings(testSettings)

        rtxDataPath = EXTENSION_FOLDER_PATH.joinpath("../../../../../data/usd/tests")
        self.open_usd_scene(rtxDataPath.joinpath("KitchenSet/Kitchen_set.usda"))

        # camera setup:
        omni.kit.commands.execute('CreatePrimWithDefaultXform',
            prim_path='/Kitchen_set/Camera',
            prim_type='Camera',
            attributes={'focusDistance': 400, 'focalLength': 24, 'clippingRange': (1, 10000000)},
            create_default_xform=False)
        omni.kit.commands.execute('TransformPrimCommand',
            path=Sdf.Path('/Kitchen_set/Camera'),
            new_transform_matrix=Gf.Matrix4d(
                0.929103525371008, 0.36981973871491086, 9.165845166192454e-16, 0.0,
                -0.07453993201486002, 0.18726775876424714, 0.9794766893921647, 0.0,
                0.3622298133483562, -0.9100352451329841, 0.201557473037752, 0.0,
                273.5074607559941, -621.0054603487383, 221.75447008156692, 1.0),
            old_transform_matrix=Gf.Matrix4d(
                0.9305632067779971, 0.36613128545789503, 9.159340032886883e-16, 0.0,
                -0.09015111941855866, 0.22912905319151347, 0.9692123877928621, 0.0,
                0.35485897742431627, -0.9019133876334855, 0.24622621174208614, 0.0,
                328.2529492634561, -760.1471818173628, 259.74075506441545, 1.0),
            time_code=Usd.TimeCode.Default(),
            had_transform_at_key=False,
            usd_context_name='')

        # remove any selection, not sure why this happens, but otherwise weird bounding
        # boxes appear
        selection = omni.usd.get_context().get_selection()
        selection.clear_selected_prim_paths()

        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, # LTC gives consistent lighting
            "/rtx/reflections/enabled" : False, # Disabling reflections.
            "/rtx/reflections/sampledLighting/enabled" : False, # LTC gives consistent lighting
            "/rtx/sceneDb/ambientLightIntensity" : float(0.0),
            "/rtx/ambientOcclusion/enabled" : False,
            "/rtx/indirectDiffuse/enabled" : False,
            "/app/viewport/grid/enabled" : False, # Disable Kit Grid
            "/app/viewport/grid/showOrigin" : False, # Disable Kit Origin
            "/app/viewport/outline/enabled" : False, # Disable Selection drawing
        }
        self.set_settings(my_settings)

    async def test_light_toggling(self):
        """
        Test RTX Light Toggling
        """

        LIGHT_VIS_PATH = ".visibility"

        # Activate the main camera:
        viewport_api = omni.kit.viewport.utility.get_active_viewport()
        viewport_api.set_active_camera("/Kitchen_set/Camera")

        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "0_all_light_sources.png")

        # Disable SphereLight
        omni.kit.commands.execute("ChangeProperty", prop_path="/SphereLight" + LIGHT_VIS_PATH, value="invisible", prev="inherited")
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "1_spherelight_off.png")

        # Disable DiskLight
        omni.kit.commands.execute("ChangeProperty", prop_path="/DiskLight" + LIGHT_VIS_PATH, value="invisible", prev="inherited")
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "2_disklight_off.png")

        # Disable DistantLight
        # Disable RectLightLight
        omni.kit.commands.execute("ChangeProperty", prop_path="/DistantLight" + LIGHT_VIS_PATH, value="invisible", prev="inherited")
        omni.kit.commands.execute("ChangeProperty", prop_path="/RectLight" + LIGHT_VIS_PATH, value="invisible", prev="inherited")
        omni.kit.commands.execute("ChangeProperty", prop_path="/DiskLight" + LIGHT_VIS_PATH, value="inherited", prev="invisible")

        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "3_two_lights_off_one_light_on.png")

        # Enable Sphere Light
        # Enable Rect Light
        # Disable Cylinder Light
        omni.kit.commands.execute("ChangeProperty", prop_path="/SphereLight" + LIGHT_VIS_PATH, value="inherited", prev="invisible")
        omni.kit.commands.execute("ChangeProperty", prop_path="/RectLight" + LIGHT_VIS_PATH, value="inherited", prev="invisible")
        omni.kit.commands.execute("ChangeProperty", prop_path="/CylinderLight" + LIGHT_VIS_PATH, value="invisible", prev="inherited")
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "4_one_light_off_two_lights_on.png")
