## Copyright (c) 2021-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 USD_DIR, EXTENSION_FOLDER_PATH
from pxr import Gf, Sdf, Usd
from pxr import UsdGeom, UsdLux

class TestRtxPostprocessingTonemapper(RtxTest):
    TEST_PATH = "tonemapping"

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

    def create_dome_light(self, name="/Xform/DomeLight"):
        omni.kit.commands.execute(
            "CreatePrim",
            prim_path=name,
            prim_type="DomeLight",
            select_new_prim=False,
            attributes={
                UsdLux.Tokens.inputsIntensity: 1,
                UsdLux.Tokens.inputsTextureFormat: UsdLux.Tokens.latlong,
                UsdLux.Tokens.inputsTextureFile: str(EXTENSION_FOLDER_PATH.joinpath("data/usd/hydra/daytime.hdr")),
                UsdGeom.Tokens.visibility: "inherited",
            } if hasattr(UsdLux.Tokens, 'inputsIntensity') else {
                UsdLux.Tokens.intensity: 1,
                UsdLux.Tokens.textureFormat: UsdLux.Tokens.latlong,
                UsdLux.Tokens.textureFile: str(EXTENSION_FOLDER_PATH.joinpath("data/usd/hydra/daytime.hdr")),
                UsdGeom.Tokens.visibility: "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()
        carb.log_info("Setting up scene for RTX postprocessing tonemapping tests.")

        # Setting that should be set before opening a new stage
        self.set_settings(testSettings)

        rtxDataPath = EXTENSION_FOLDER_PATH.joinpath("../../../../../data/usd/tests")
        self.open_usd_scene(rtxDataPath.joinpath("BallCluster/ballcluster_stage.usda"))

        # Settings that should be set after opening a new stage
        self.set_settings(postLoadTestSettings)

        # camera setup:
        omni.kit.commands.execute('CreatePrimWithDefaultXform',
            prim_path='/Ballcluster_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('/Ballcluster_set/Camera'),
            new_transform_matrix=Gf.Matrix4d(-0.6550639249522971, 0.7555734605093614, 6.112664085837494e-16, 0.0,
                    0.008064090705444693, 0.006991371699474781, 0.9999430439594318, 0.0,
                    0.7555304260366915, 0.6550266151048126, -0.010672808306432642, 0.0,
                    543.7716131933593, 536.4127523809548, 92.10006955095595, 1.0),
            old_transform_matrix=Gf.Matrix4d(-0.7071067811865474, 0.7071067811865478, 5.551115215516806e-17, 0.0,
               -0.4082482839677534, -0.4082482839677533, 0.8164965874238357, 0.0,
               0.5773502737830692, 0.5773502737830688, 0.5773502600027394, 0.0,
               500.0, 500.0, 500.0, 1.0),
            time_code=Usd.TimeCode.Default(),
            had_transform_at_key=False,
            usd_context_name='')


        # Create a domelight to get some high dynamic range going:
        self.domelight_prim = self.create_dome_light()
        # Disable this object so that the dome light is actually visible:
        omni.kit.commands.execute('DeletePrims', paths=['/Xform/sky_sphere_emiss'], destructive=False)
        await wait_for_update()

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

        await omni.kit.app.get_app().next_update_async()

        my_settings = {
            "/app/hydraEngine/waitIdle" : True,
            "/app/renderer/waitIdle" : True,
            "/app/asyncRenderingLowLatency" : False,
            "/rtx-transient/resourcemanager/genMipsForNormalMaps" : False,
            "/rtx-transient/resourcemanager/texturestreaming/async" : False,
            "/rtx-transient/samplerFeedbackTileSize" : 1,
            "/rtx/hydra/perMaterialSyncLoads" : True,
            "/rtx/post/aa/op" : 0, # 0 = None, 2 = FXAA
            "/renderer/multiGpu/maxGpuCount" : 1,
            "/rtx/gatherColorToDisplayDevice" : True,

            "/rtx/rendermode" : 'PathTracing',
            "/rtx/raytracing/lightcache/spatialCache/enabled" : False,
            "/rtx/pathtracing/lightcache/cached/enabled" : False,
            "/rtx/pathtracing/cached/enabled" : False,
            "/rtx/pathtracing/optixDenoiser/enabled" : False,
            "/rtx/pathtracing/spp" : 1,
            "/rtx/pathtracing/totalSpp" : 1,
            "/rtx/pathtracing/maxBounces" : 2,

            "/rtx/hydra/perMaterialSyncLoads" : True,
        }

        self.set_settings(my_settings)
        await wait_for_update()

    # Marking unstable due to OM-86613
    async def test_UNSTABLE_postprocessing_tonemappers_UNSTABLE(self):
        """
        Test RTX Postprocessing Tonemappers
        """

        self.set_settings({"/rtx/post/tonemap/op" : "0"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "0_tonemapper_test_clamp.png")

        # Make the domelight 'visible' for all non-clamp tonemappers
        # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
        if hasattr(UsdLux.Tokens, 'inputsIntensity'):
            self.domelight_prim.GetAttribute(UsdLux.Tokens.inputsIntensity).Set(200)
        else:
            self.domelight_prim.GetAttribute(UsdLux.Tokens.intensity).Set(200)

        self.set_settings({"/rtx/post/tonemap/op" : "1"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "1_tonemapper_test_linear.png")

        self.set_settings({"/rtx/post/tonemap/op" : "2"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "2_tonemapper_test_reinhard.png")

        self.set_settings({"/rtx/post/tonemap/op" : "3"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "3_tonemapper_test_reinhard_modified.png")

        self.set_settings({"/rtx/post/tonemap/op" : "4"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "4_tonemapper_test_hejlhablealu.png")

        self.set_settings({"/rtx/post/tonemap/op" : "5"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "5_tonemapper_test_hableuc2.png")

        self.set_settings({"/rtx/post/tonemap/op" : "6"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "6_tonemapper_test_aces.png")

        self.set_settings({"/rtx/post/tonemap/op" : "7"})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "7_tonemapper_test_iray.png")
