## 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.test
import omni.ui as ui
from omni.kit import ui_test
from omni.kit.test_suite.helpers import wait_stage_loading
from omni.ui.tests.test_base import OmniUiTest
from pxr import Kind, Sdf, Gf, UsdGeom
from pathlib import Path
import carb
from carb.input import KeyboardInput

def equal_eps(a, b, esp = 0.0001):
    return abs(a - b) < esp

def equal_scale_eps(t, s):
    m = Gf.Vec3d(0.0)
    m[0] = t.GetRow(0).GetLength()
    m[1] = t.GetRow(1).GetLength()
    m[2] = t.GetRow(2).GetLength()

    #print (m)

    return equal_eps(m[0], s[0]) and equal_eps(m[1], s[1]) and equal_eps(m[2], s[2])

class TestTransformWidget(OmniUiTest):
    # Before running each test
    async def setUp(self):
        await super().setUp()

        extension_root_folder = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))

        self._golden_img_dir = extension_root_folder.joinpath("data/test_data/golden_img")
        self._usd_path = extension_root_folder.joinpath("data/test_data/test_map")

        from omni.kit.property.usd.usd_attribute_widget import UsdPropertiesWidget
        import omni.kit.window.property as p

        self._w = p.get_window()

    # After running each test
    async def tearDown(self):
        await super().tearDown()

    async def test_transform_property_srt(self):
        usd_context = omni.usd.get_context()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=250,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_srt.png")

    async def test_transform_property_srt_filter(self):
        usd_context = omni.usd.get_context()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=250,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)

        try:
            self._w._searchfield._search_field.model.as_string = "tr"
            self._w._searchfield._set_in_searching(True)

            # Need to wait for an additional frames for omni.ui rebuild to take effect
            await ui_test.human_delay(10)

            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_srt_filter.png")
        finally:
            self._w._searchfield._search_field.model.as_string = ""
            self._w._searchfield._set_in_searching(False)

    async def test_transform_property_sqt(self):
        usd_context = omni.usd.get_context()

        # golden image expects the UI to be switched to "orient as rotate"
        settings = carb.settings.get_settings()
        doar_path = "/persistent/app/uiSettings/DisplayOrientAsRotate"
        doar_val = settings.get_as_bool(doar_path)
        settings.set(doar_path, True)

        try:
            await self.docked_test_window(
                window=self._w._window,
                width=450,
                height=250,
                restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
                restore_position = ui.DockPosition.BOTTOM)

            test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
            await usd_context.open_stage_async(str(test_file_path))
            await wait_stage_loading()

            # Select the prim.
            usd_context.get_selection().set_selected_prim_paths(["/World/Capsule"], True)

            # Need to wait for an additional frames for omni.ui rebuild to take effect
            await ui_test.human_delay(10)
        finally:
            carb.settings.get_settings().set(doar_path, doar_val)
            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_sqt.png")

    async def test_transform_property_sqt2(self):
        usd_context = omni.usd.get_context()

        # golden image expects the UI to be switched to "orient as rotate"
        settings = carb.settings.get_settings()
        doar_path = "/persistent/app/uiSettings/DisplayOrientAsRotate"
        doar_val = settings.get_as_bool(doar_path)
        settings.set(doar_path, True)

        try:
            await self.docked_test_window(
                window=self._w._window,
                width=450,
                height=250,
                restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
                restore_position = ui.DockPosition.BOTTOM,
                block_devices=False)

            test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
            await usd_context.open_stage_async(str(test_file_path))
            await wait_stage_loading()

            # Select the prim.
            usd_context.get_selection().set_selected_prim_paths(["/World/Capsule"], True)

            # Need to wait for an additional frames for omni.ui rebuild to take effect
            await ui_test.human_delay(10)

            #now find a button for orient and click it to swap to test swapping
            widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Button[*].text=='Orient'")
            await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10))

        finally:
            await self.capture_and_compare(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_sqt2.png")

            # add another test to click it back
            widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Button[*].text=='Orient'")
            await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10))

            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_sqt3.png")

            carb.settings.get_settings().set(doar_path, doar_val)

    async def test_transform_property_matrix(self):
        usd_context = omni.usd.get_context()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=250,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cylinder"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_matrix.png")

    async def test_transform_property_pivot(self):
        usd_context = omni.usd.get_context()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=280,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cone"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_pivot.png")

    async def test_transform_property_disabled_translate_op(self):
        usd_context = omni.usd.get_context()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=260,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Sphere"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_disabled_translate_op.png")

    async def test_transform_property_euler_to_quat_sync(self):
        usd_context = omni.usd.get_context()
        usd_context.new_stage()

        settings = carb.settings.get_settings()
        doar_path = "/persistent/app/uiSettings/DisplayOrientAsRotate"
        doar_val = settings.get_as_bool(doar_path)
        carb.settings.get_settings().set(doar_path, True)

        try:
            await self.docked_test_window(
                window=self._w._window,
                width=640,
                height=480,
                restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
                restore_position = ui.DockPosition.BOTTOM)

            # we actually want mouse and keyboard to be active
            omni.appwindow.get_default_app_window().set_input_blocking_state(carb.input.DeviceType.MOUSE, False)
            omni.appwindow.get_default_app_window().set_input_blocking_state(carb.input.DeviceType.KEYBOARD, False)

            await ui_test.wait_n_updates(10)

            path = "/Cube"
            cube_geom = UsdGeom.Cube.Define(usd_context.get_stage(), path)
            cube_geom.AddOrientOp().Set(Gf.Quatf(1.0))

            usd_context.get_selection().set_selected_prim_paths([path], False)
            await ui_test.find("Property").focus()

            transform_euler = ui_test.find("Property//Frame/**/*.identifier=='euler_xformOp:orient'")

            # change Y axis to 180 digit by digit as a user would do
            await transform_euler.double_click(human_delay_speed=10)
            for s in "180":
                await ui_test.emulate_char_press(s)
                await ui_test.wait_n_updates(10)
            await ui_test.emulate_keyboard_press(carb.input.KeyboardInput.ENTER)
        finally:
            await self.finalize_test_no_image()
            carb.settings.get_settings().set(doar_path, doar_val)

        orient = Gf.Quatf(1.0)
        for op in cube_geom.GetOrderedXformOps():
            if op.GetOpType() == UsdGeom.XformOp.TypeOrient:
                orient = op.Get()
        self.assertTrue(orient.imaginary[1] != 0)

    async def test_transform_property_multi(self):
        usd_context = omni.usd.get_context()

        # golden image expects the UI to be switched to "orient as rotate"
        try:
            await self.docked_test_window(
                window=self._w._window,
                width=450,
                height=250,
                restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
                restore_position = ui.DockPosition.BOTTOM)

            test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
            await usd_context.open_stage_async(str(test_file_path))
            await wait_stage_loading()

            # Select the prim.
            usd_context.get_selection().set_selected_prim_paths(["/World/Capsule", "/World/Cube", "/World/Cylinder"], True)

            # Need to wait for an additional frames for omni.ui rebuild to take effect
            await ui_test.human_delay(10)

        finally:
            await self.capture_and_compare(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_multi.png")

            # Select the prim.
            usd_context.get_selection().set_selected_prim_paths(["/World/Capsule", "/World/Cube"], True)

            # Need to wait for an additional frames for omni.ui rebuild to take effect
            await ui_test.human_delay(10)

            #text of mixed is moving and it's causing very unstable result for testing, so increasing threshold to be 2.
            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_multi2.png", threshold=2)

    async def test_transform_scale(self):
        usd_context = omni.usd.get_context()

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        stage = usd_context.get_stage()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)
        prim = stage.GetPrimAtPath("/World/Cube")
        # Let UI build
        # (Note: I'm not sure why we need to wait more than a few frames, but empirically
        # this seems to be the case. So just wait 20 frames to make sure UI is there.)
        await ui_test.wait_n_updates(20)

        # Now we're in offset mode.
        stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='scale_t_stack'")
        self.assertIsNotNone(stack)
        # Here's a funny thing -- the double-click clicks the center of the ui element by default.
        # That amounts to setting offsets in Y.
        # Rather than work around that, we're just going to use the Y component for our offset.

        # XXX for some reason, using field.input doesn't give the desired results here --
        # the field's text isn't selected after double-clicking, so we end up entering
        # "0.0100" instead of "100".
        # As a workaround, we'll just break input down into its components (double-click, enter characters,
        # hit enter), with the addition of three backspaces to clear the field's default "0.0"
        await stack.click()
        await stack.double_click()
        await ui_test.human_delay(10)
        for _ in range(3):
            await ui_test.emulate_keyboard_press(KeyboardInput.BACKSPACE)
        await ui_test.emulate_char_press("2")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_scale_eps(transform, Gf.Vec3d(1.0, 2.0, 1.0)))

        # Apply another offset, then check the prim's translation
        await stack.double_click()
        await ui_test.human_delay(10)
        for _ in range(3):
            await ui_test.emulate_keyboard_press(KeyboardInput.BACKSPACE)
        await ui_test.emulate_char_press("*2")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_scale_eps(transform, Gf.Vec3d(1.0, 4.0, 1.0)))

        # Click and drag field, the check the prim's translation
        drag_vector = stack.center
        drag_vector.x = drag_vector.x + 30
        await ui_test.human_delay(30)
        await ui_test.emulate_mouse_drag_and_drop(stack.center, drag_vector)
        await ui_test.wait_n_updates(2)

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_scale_eps(transform, Gf.Vec3d(1.0, 7.0, 1.0)))

        await self.finalize_test_no_image()

        await ui_test.wait_n_updates(20)

    async def test_transform_toggle_scale(self):
        usd_context = omni.usd.get_context()

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        stage = usd_context.get_stage()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)
        prim = stage.GetPrimAtPath("/World/Cube")

        # Let UI build
        # (Note: I'm not sure why we need to wait more than a few frames, but empirically
        # this seems to be the case. So just wait 20 frames to make sure UI is there.)
        await ui_test.wait_n_updates(20)

        button = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Button[*].identifier=='toggle_link_scale'")

        # We should be able to find the mode toggle button,
        # but until we're in offset mode we should not have any offset fields
        self.assertIsNotNone(button)
        #enable toggle mode
        await button.click(pos=button.position+ui_test.Vec2(button.widget.computed_content_width/2, 5))
        await ui_test.wait_n_updates(20)

        # Now we're in offset mode.
        stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='scale_t_stack'")
        self.assertIsNotNone(stack)
        # Here's a funny thing -- the double-click clicks the center of the ui element by default.
        # That amounts to setting offsets in Y.
        # Rather than work around that, we're just going to use the Y component for our offset.

        # XXX for some reason, using field.input doesn't give the desired results here --
        # the field's text isn't selected after double-clicking, so we end up entering
        # "0.0100" instead of "100".
        # As a workaround, we'll just break input down into its components (double-click, enter characters,
        # hit enter), with the addition of three backspaces to clear the field's default "0.0"
        await stack.click()
        await stack.double_click()
        await ui_test.human_delay(10)
        for _ in range(3):
            await ui_test.emulate_keyboard_press(KeyboardInput.BACKSPACE)
        await ui_test.emulate_char_press("2")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_scale_eps(transform, Gf.Vec3d(2.0, 2.0, 2.0)))

        await self.finalize_test_no_image()

        await ui_test.wait_n_updates(20)

    async def test_transform_reset_stack(self):
        usd_context = omni.usd.get_context()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=280,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        stage = omni.usd.get_context().get_stage()
        prim = stage.GetPrimAtPath("/World/Cube")

        # Get the XformCommonAPI
        xform_api = UsdGeom.XformCommonAPI(prim)

        #Set the ResetXformStack. UI should responds accordingly
        xform_api.SetResetXformStack(True)

        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_reset.png")

    async def test_transform_rotate_scalar(self):
        usd_context = omni.usd.get_context()

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        stage = usd_context.get_stage()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube_2"], True)
        prim = stage.GetPrimAtPath("/World/Cube_2")
        # Let UI build
        # (Note: I'm not sure why we need to wait more than a few frames, but empirically
        # this seems to be the case. So just wait 20 frames to make sure UI is there.)
        await ui_test.wait_n_updates(20)

        # Now we're in offset mode.
        x_stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='rotate_scalar_stack_x'")
        self.assertIsNotNone(x_stack)
        # Here's a funny thing -- the double-click clicks the center of the ui element by default.
        # That amounts to setting offsets in Y.
        # Rather than work around that, we're just going to use the Y component for our offset.

        rotate_attr = prim.GetAttribute("xformOp:rotateX")
        self.assertIsNotNone(rotate_attr)
        self.assertTrue(rotate_attr.Get() == -30.0)

        # XXX for some reason, using field.input doesn't give the desired results here --
        # the field's text isn't selected after double-clicking, so we end up entering
        # "0.0100" instead of "100".
        # As a workaround, we'll just break input down into its components (double-click, enter characters,
        # hit enter), with the addition of three backspaces to clear the field's default "0.0"
        await x_stack.click()
        await x_stack.double_click()
        await ui_test.human_delay(10)
        for _ in range(5):
            await ui_test.emulate_keyboard_press(KeyboardInput.BACKSPACE)
        await ui_test.emulate_char_press("20")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        self.assertTrue(rotate_attr.Get() == 20.0)

        y_stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='rotate_scalar_stack_y'")
        self.assertIsNotNone(y_stack)
        # Here's a funny thing -- the double-click clicks the center of the ui element by default.
        # That amounts to setting offsets in Y.
        # Rather than work around that, we're just going to use the Y component for our offset.

        rotate_attr = prim.GetAttribute("xformOp:rotateY")
        self.assertIsNotNone(rotate_attr)
        self.assertTrue(rotate_attr.Get() == 0.0)

        # XXX for some reason, using field.input doesn't give the desired results here --
        # the field's text isn't selected after double-clicking, so we end up entering
        # "0.0100" instead of "100".
        # As a workaround, we'll just break input down into its components (double-click, enter characters,
        # hit enter), with the addition of three backspaces to clear the field's default "0.0"
        await y_stack.click()
        await y_stack.double_click()
        await ui_test.human_delay(10)
        for _ in range(5):
            await ui_test.emulate_keyboard_press(KeyboardInput.BACKSPACE)
        await ui_test.emulate_char_press("30")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        self.assertTrue(rotate_attr.Get() == 30.0)
        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_eps(transform.ExtractRotation().GetAngle(), 35.92772))
        #test transform
        await self.finalize_test_no_image()

        await ui_test.wait_n_updates(20)
    '''
    #this test doesn't work because the kit has a way to enter xformOp:translate to trigger
    #build widget, where test won't be able to do this. OM-112507

    #https://nvidia.slack.com/archives/CURCH7KU2/p1697658671468429?thread_ts=1697204512.040859&cid=CURCH7KU2

    async def test_transform_property_wgs84(self):
        usd_context = omni.usd.get_context()
        await self.docked_test_window(
            window=self._w._window,
            width=640,
            height=480,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        test_file_path = self._usd_path.joinpath("wgs84/deutschebahn-rails.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/RootGeoReference/World/CurveXform26Geo"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        try:
            self._w._searchfield._search_field.model.as_string = "tr"
            self._w._searchfield._set_in_searching(True)

            # Need to wait for an additional frames for omni.ui rebuild to take effect
            await ui_test.human_delay(10)

            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_transform_property_wgs84.png")
        finally:
            self._w._searchfield._search_field.model.as_string = ""
            self._w._searchfield._set_in_searching(False)
    '''