## 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.appwindow
import omni.kit.app
import omni.kit.test
import omni.kit.commands
import omni.timeline
import omni.usd
import carb
import carb.settings
import carb.windowing
import inspect
import pathlib
from omni.kit.test_helpers_gfx.compare_utils import finalize_capture_and_compare, ComparisonMetric
from omni.kit.viewport.utility import get_active_viewport_window, next_viewport_frame_async
from omni.kit.viewport.utility.tests.capture import capture_viewport_and_wait
from pxr import Sdf, Gf, Usd, UsdGeom, UsdLux

# This settings should be set before stage opening/creation
testSettings = {
    "/app/window/hideUi": True,
    "/app/asyncRendering": False,
    "/app/docks/disabled": True,
    "/app/window/scaleToMonitor": False,
    "/app/viewport/forceHideFps": True,
    "/app/captureFrame/setAlphaTo1": True,
    "/rtx/materialDb/syncLoads": True,
    "/omni.kit.plugin/syncUsdLoads": True,
    "/rtx/hydra/materialSyncLoads": True,
    "/renderer/multiGpu/autoEnable": False,
    "/persistent/app/viewport/displayOptions": 0,
    "/app/viewport/grid/enabled": False,
    "/app/viewport/show/lights": False,
    "/persistent/app/primCreation/PrimCreationWithDefaultXformOps": True,
    "/rtx-transient/resourcemanager/texturestreaming/async": False,
    "/app/viewport/outline/enabled": True
}

# Settings that should override settings that was set on stage opening/creation
postLoadTestSettings = {
    "/rtx/post/aa/op": 0,
    "/rtx/pathtracing/lightcache/cached/enabled": False,
    "/rtx/raytracing/lightcache/spatialCache/enabled": False,
    "/rtx/sceneDb/ambientLightIntensity": 1.0,
    "/rtx/indirectDiffuse/enabled": False,
}

postLoadSkelTestSettings = {
    "/rtx/post/aa/op": 0,
    "/rtx/shadows/enabled": False,
    "/rtx/reflections/enabled": False,
    "/rtx/ambientOcclusion/enabled": False,
    "/rtx/post/tonemap/op": 1,
    "/renderer/multiGpu/autoEnable": False,
}

RENDER_WIDTH_SETTING = "/app/renderer/resolution/width"
RENDER_HEIGHT_SETTING = "/app/renderer/resolution/height"

OUTPUTS_DIR = pathlib.Path(omni.kit.test.get_test_output_path())

EXTENSION_FOLDER_PATH = pathlib.Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))
GOLDEN_DIR = EXTENSION_FOLDER_PATH.joinpath("data/golden")
USD_DIR = EXTENSION_FOLDER_PATH.joinpath("data/usd")
VOLUMES_DIR = EXTENSION_FOLDER_PATH.joinpath("data/volumes")


async def next_resize_async():
    """
    Wait for the next event in the resize event stream of IAppWindow::getWindowResizeEventStream.

    We need it because the window resize event stream is independent of IApp::getUpdateEventStream. Without this
    function it's possible that resize happens several updates after. It's reproducable on Linux Release build.
    """
    return await omni.appwindow.get_default_app_window().get_window_resize_event_stream().next_event()


async def wait_for_update(usd_context=omni.usd.get_context(), wait_frames=10):
    max_loops = 0
    while max_loops < wait_frames:
        _, files_loaded, total_files = usd_context.get_stage_loading_status()
        await omni.kit.app.get_app().next_update_async()
        if files_loaded or total_files:
            continue
        max_loops = max_loops + 1


def set_cursor_position(pos):
    app_window = omni.appwindow.get_default_app_window()
    windowing = carb.windowing.acquire_windowing_interface()
    os_window = app_window.get_window()
    windowing.set_cursor_position(os_window, pos)


def set_transform_helper(
    prim_path,
    translate=Gf.Vec3d(0, 0, 0),
    euler=Gf.Vec3d(0, 0, 0),
    scale=Gf.Vec3d(1, 1, 1),
):
    rotation = (
        Gf.Rotation(Gf.Vec3d.ZAxis(), euler[2])
        * Gf.Rotation(Gf.Vec3d.YAxis(), euler[1])
        * Gf.Rotation(Gf.Vec3d.XAxis(), euler[0])
    )
    xform = Gf.Matrix4d().SetScale(scale) * Gf.Matrix4d().SetRotate(rotation) * Gf.Matrix4d().SetTranslate(translate)
    omni.kit.commands.execute(
        "TransformPrimCommand",
        path=prim_path,
        new_transform_matrix=xform,
    )

async def setup_viewport_test_window(resolution_x: int, resolution_y: int, position_x: int = 0, position_y: int = 0):
    from omni.kit.viewport.utility import get_active_viewport_window
    viewport_window = get_active_viewport_window()
    if viewport_window:
        viewport_window.position_x = position_x
        viewport_window.position_y = position_y
        viewport_window.width = resolution_x
        viewport_window.height = resolution_y
        viewport_window.viewport_api.resolution = (resolution_x, resolution_y)
    return viewport_window


class RtxTest(omni.kit.test.AsyncTestCase):
    THRESHOLD = 1e-5
    WINDOW_SIZE = (640, 480)

    def __init__(self, tests=()):
        super().__init__(tests)
        self._saved_width = None
        self._saved_height = None
        self._savedSettings = {}
        self._failedImages = []

    @property
    def __test_name(self) -> str:
        """
        The full name of the test.

        It has the name of the module, class and the current test function. We use the stack to get the name of the test
        function and since it's only called from create_test_window and finalize_test, we get the third member.
        """
        return f"{self.__module__}.{self.__class__.__name__}.{inspect.stack()[2][3]}"

    async def create_test_area(self, width: int = 256, height: int = 256):
        """Resize the main window"""
        app_window = omni.appwindow.get_default_app_window()

        await omni.usd.get_context().new_stage_async()
        viewport_window = await setup_viewport_test_window(width, height)
        self.assertTrue(viewport_window is not None, "No active viewport window found.")

        # Current main window size
        current_width = app_window.get_width()
        current_height = app_window.get_height()

        # If the main window is already has requested size, do nothing
        if width == current_width and height == current_height:
            self._saved_width = None
            self._saved_height = None
        else:
            # Save the size of the main window to be able to restore it at the end of the test
            self._saved_width = current_width
            self._saved_height = current_height

            app_window.resize(width, height)

            # Wait for getWindowResizeEventStream
            await next_resize_async()

        # Wait until the Viewport has delivered some frames
        await next_viewport_frame_async(viewport_window.viewport_api, 0)

    async def screenshot_and_diff(self, golden_img_dir: pathlib.Path, output_subdir=None,
                                  golden_img_name=None, threshold=THRESHOLD):
        """
        Capture the current frame and compare it with the golden image. Assert if the diff is more than given threshold.

        This method differs from capture_and_compare in that it lets callers outside omni.rtx.tests
        capture and compare images in their own directories. The screen captures will be placed in a
        common place with the rtx.tests output, either in a subdirectory passed in as output_subdir,
        or in a directory named for the test module. Golden images will be found in the directory passed
        in as golden_img_dir, this is the only required parameter.
        """
        if not golden_img_dir:
            self.assertTrue(golden_img_dir, "A valid golden image dir is a required parameter")

        if not output_subdir:
            output_subdir = f"{self.__module__}"

        output_img_dir = OUTPUTS_DIR.joinpath(output_subdir)

        if not golden_img_name:
            golden_img_name = f"{self.__test_name}.png"

        return await self._capture_and_compare(golden_img_name, threshold, output_img_dir, golden_img_dir)

    async def capture_and_compare(self, img_subdir: pathlib.Path = None, golden_img_name=None, threshold=THRESHOLD,
                                  metric: ComparisonMetric = ComparisonMetric.MEAN_ERROR_SQUARED):
        """
        Capture current frame and compare it with the golden image. Assert if the diff is more than given threshold.
        """
        golden_img_dir = GOLDEN_DIR.joinpath(img_subdir)
        output_img_dir = OUTPUTS_DIR.joinpath(img_subdir)

        if not golden_img_name:
            golden_img_name = f"{self.__test_name}.png"

        return await self._capture_and_compare(golden_img_name, threshold, output_img_dir, golden_img_dir, metric)

    async def _capture_and_compare(self, golden_img_name, threshold, output_img_dir: pathlib.Path,
                                   golden_img_dir: pathlib.Path,
                                   metric: ComparisonMetric = ComparisonMetric.MEAN_ERROR_SQUARED):

        # Capture directly from the Viewport's texture, not from UI swapchain
        await capture_viewport_and_wait(golden_img_name, output_img_dir)

        # Do the image comparison now
        diff = finalize_capture_and_compare(golden_img_name, threshold, output_img_dir, golden_img_dir, metric=metric)

        if diff != 0:
            carb.log_warn(f"[{self.__test_name}] the generated image {golden_img_name} has max difference {diff}")

        if (diff is not None) and diff >= threshold:
            self._failedImages.append(golden_img_name)
        return diff

    def add_dir_light(self):
        omni.kit.commands.execute(
            "CreatePrimCommand",
            prim_path="/World/Light",
            prim_type="DistantLight",
            select_new_prim=False,
            # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
            attributes={UsdLux.Tokens.inputsAngle: 1.0, UsdLux.Tokens.inputsIntensity: 3000} if hasattr(UsdLux.Tokens, 'inputsIntensity') else
                {UsdLux.Tokens.angle: 1.0, UsdLux.Tokens.intensity: 3000},
            create_default_xform=True,
        )

    def add_floor(self):
        floor_path = "/World/Floor"
        omni.kit.commands.execute(
            "CreatePrimCommand",
            prim_path=floor_path,
            prim_type="Cube",
            select_new_prim=False,
            attributes={UsdGeom.Tokens.size: 100},
        )
        floor_prim = self.ctx.get_stage().GetPrimAtPath(floor_path)
        floor_prim.CreateAttribute("xformOp:scale", Sdf.ValueTypeNames.Double3, False).Set(Gf.Vec3d(25, 0.1, 25))
        floor_prim.CreateAttribute("xformOpOrder", Sdf.ValueTypeNames.String, False).Set(["xformOp:scale"])

    def open_usd(self, usdSubpath: pathlib.Path):
        path = USD_DIR.joinpath(usdSubpath)
        omni.usd.get_context().open_stage(str(path))

    def get_volumes_path(self, volumeSubPath: pathlib.Path):
        return Sdf.AssetPath(str(VOLUMES_DIR.joinpath(volumeSubPath)))

    def set_settings(self, newSettings):
        settingsAPI = carb.settings.get_settings()
        for s, v in newSettings.items():
            if s not in self._savedSettings:  # Remember old setting only when it was changed first time
                self._savedSettings[s] = settingsAPI.get(s)
            if v is None:  # hideUi sometimes is None and it hangs kit
                v = False
            settingsAPI.set(s, v)

    def set_camera(self, cameraPos=None, targetPos=None):
        from omni.kit.viewport.utility.camera_state import ViewportCameraState
        camera_state = ViewportCameraState("/OmniverseKit_Persp")
        camera_state.set_position_world(cameraPos, True)
        camera_state.set_target_world(targetPos, True)

    async def setUp_internal(self):
        self.ctx = omni.usd.get_context()
        await self.create_test_area(self.WINDOW_SIZE[0], self.WINDOW_SIZE[1])

    async def setUp(self):
        await self.setUp_internal()

    async def tearDown(self):
        # Restore main window resolution if it was saved
        if self._saved_width is not None and self._saved_height is not None:
            app_window = omni.appwindow.get_default_app_window()
            app_window.resize(self._saved_width, self._saved_height)
            # Wait for getWindowResizeEventStream
            await next_resize_async()
        self.set_settings(self._savedSettings)
        self.ctx.close_stage()
        for imgName in self._failedImages:
            carb.log_warn(f"[{self.__test_name}] The image {imgName} doesn't match the golden")
        hasFailed = len(self._failedImages)
        self._failedImages = []
        self.assertEqual(hasFailed, 0)
