## 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.
##

from carb.input import KeyboardInput
import omni.kit.test
import omni.kit.ui_test as ui_test
import omni.ui as ui
import omni.usd
from omni.ui.tests.test_base import OmniUiTest
from pxr import Gf

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

class TestTransformOffset(OmniUiTest):
    async def setUp(self):
        await super().setUp()

        from omni.kit.property.transform.scripts.transform_properties import TEST_DATA_PATH
        self._usd_path = TEST_DATA_PATH.absolute()

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

        self._w = p.get_window()

    async def tearDown(self):
        await super().tearDown()

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

        test_file_path = self._usd_path.joinpath("offset_test.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()

        self.assertIsNotNone(stage)

        prim = stage.GetPrimAtPath(f"{stage.GetDefaultPrim().GetPath()}/Cube")
        self.assertTrue(prim.IsValid())

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)

        # 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/**/Button[*].identifier=='offset_mode_toggle'")

        # 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)

        await button.click()
        await ui_test.wait_n_updates(20)

        # Now we're in offset mode.
        stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='rotate_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("10")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_eps(transform.ExtractRotation().GetAngle(), 10.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_eps(transform.ExtractRotation().GetAngle(), 20.0))

        # Click and drag field, the check the prim's translation
        drag_vector = stack.center
        drag_vector.x = drag_vector.x + 100
        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_eps(transform.ExtractRotation().GetAngle(), 21.0))

        await self.finalize_test_no_image()

        # Finally, toggle offset mode again so subsequent tests will not start in offset mode
        await button.click()
        await ui_test.wait_n_updates(20)


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

        test_file_path = self._usd_path.joinpath("offset_test.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()

        self.assertIsNotNone(stage)

        prim = stage.GetPrimAtPath(f"{stage.GetDefaultPrim().GetPath()}/Cube")
        self.assertTrue(prim.IsValid())

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)

        # 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/**/Button[*].identifier=='offset_mode_toggle'")
        stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='translate_stack'")

        # 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)
        self.assertIsNone(stack)

        await button.click()
        await ui_test.wait_n_updates(20)

        # Now we're in offset mode.
        stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='translate_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("100")
        await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(transform.ExtractTranslation() == Gf.Vec3d(0.0, 100.0, 0.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(transform.ExtractTranslation() == Gf.Vec3d(0.0, 200.0, 0.0))

        # Click and drag field, the check the prim's translation
        drag_vector = stack.center
        drag_vector.x = drag_vector.x + 100
        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(transform.ExtractTranslation() == Gf.Vec3d(0.0, 201.0, 0.0))

        await self.finalize_test_no_image()

        # Finally, toggle offset mode again so subsequent tests will not start in offset mode
        await button.click()
        await ui_test.wait_n_updates(20)


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

        test_file_path = self._usd_path.joinpath("offset_test.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()

        self.assertIsNotNone(stage)

        prim = stage.GetPrimAtPath(f"{stage.GetDefaultPrim().GetPath()}/Cube")
        self.assertTrue(prim.IsValid())

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)

        # 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/**/Button[*].identifier=='offset_mode_toggle'")

        # 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)

        await button.click()
        await ui_test.wait_n_updates(20)

        # Now we're in offset mode.
        stack = ui_test.find("Property//Frame/**/HStack[*].identifier=='scale_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)

        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])

        transform = omni.usd.get_world_transform_matrix(prim)
        self.assertTrue(equal_scale_eps(transform, Gf.Vec3d(1.0, 3.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, 6.0, 1.0)))

        # Click and drag field, the check the prim's translation
        drag_vector = stack.center
        drag_vector.x = drag_vector.x + 100
        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()

        # Finally, toggle offset mode again so subsequent tests will not start in offset mode
        await button.click()
        await ui_test.wait_n_updates(20)
