## 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 weakref
import omni.kit.app
import omni.kit.commands
import omni.kit.test
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 pathlib import Path
from pxr import UsdGeom, Sdf, Gf
from omni.kit.property.usd.prim_selection_payload import PrimSelectionPayload
from ..scripts.xform_op_utils import get_op_type_name, _add_trs_op, get_op_name_suffix, get_op_precision, get_inverse_op_Name, RESET_XFORM_STACK, get_op_attr_name
import carb
import carb.settings

class TestTransformCommands(OmniUiTest):

    xform_op_names = ["xformOp:translate", "xformOp:rotateXYZ", "xformOp:orient", "xformOp:scale", "xformOp:translate:pivot", "xformOp:transform"]
    prim_paths =["/World/Capsule", "/World/Cube", "/World/Cylinder", "/World/Cone"]
    # 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._usd_path = extension_root_folder.joinpath("data/test_data/test_map")

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

    def include_op_code(self, xform : UsdGeom.Xformable, op_name):
        return self.get_index_op_code(xform, op_name) != -1

    def get_index_op_code(self, xform : UsdGeom.Xformable, op_name):
        xform_ops = xform.GetOrderedXformOps()
        for i, op in enumerate(xform_ops):
            if op.GetOpName() == op_name:
                return i

        return -1

    async def add_xform_ops_test(self, stage, prim_path, precision, undo_after=True):
        prim = stage.GetObjectAtPath(prim_path)
        xform = UsdGeom.Xformable(prim)

        usd_context = omni.usd.get_context()
        usd_context.get_selection().set_selected_prim_paths([prim_path], True)

        await ui_test.human_delay(3)

        ops_exists = []
        attr_exists = []

        for op_name in self.xform_op_names:
            ops_exists.append(self.include_op_code(xform, op_name))
            attr_exists.append(stage.GetObjectAtPath(prim_path + "." + op_name) != None)

        payload = PrimSelectionPayload( weakref.ref(stage), [prim.GetPath()])
        omni.kit.commands.execute('AddXformOp',
            payload=payload,
            precision=precision,
            rotation_order='XYZ',
            add_translate_op=True,
            add_rotateXYZ_op=True,
            add_orient_op=True,
            add_scale_op=True,
            add_transform_op=True,
            add_pivot_op=True)

        await ui_test.human_delay(3)

        xform_ops = xform.GetOrderedXformOps()

        for op_name in self.xform_op_names:
            #make sure rotation order XYZ ops_exists
            self.assertTrue(self.include_op_code(xform, op_name))
            attr = stage.GetObjectAtPath(prim_path + "." + op_name)
            self.assertTrue(attr)

        if undo_after:
            #undo
            omni.kit.commands.execute("Undo")
            await ui_test.human_delay(3)

            for i, op_name in enumerate(self.xform_op_names):
                if not ops_exists[i]:
                    self.assertTrue(not self.include_op_code(xform, op_name))

    async def remove_xform_ops_test(self, stage, prim_path):
        prim = stage.GetObjectAtPath(prim_path)
        xform = UsdGeom.Xformable(prim)

        usd_context = omni.usd.get_context()
        usd_context.get_selection().set_selected_prim_paths([prim_path], True)
        await ui_test.human_delay(3)

        for op_name in self.xform_op_names:
            index = self.get_index_op_code(xform, op_name)
            omni.kit.commands.execute('RemoveXformOpAndAttrbute',
                op_order_attr_path=Sdf.Path(prim_path + '.xformOpOrder'),
                op_name=op_name,
                op_order_index=index)

            await ui_test.human_delay(3)
            self.assertTrue(not self.include_op_code(xform, op_name))

        for op_name in self.xform_op_names:
            omni.kit.commands.execute("Undo")

        for op_name in self.xform_op_names:
            self.assertTrue(self.include_op_code(xform, op_name))

    async def test_transform_add_remove_commands(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 wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()

        precision_types = [UsdGeom.XformOp.PrecisionDouble, UsdGeom.XformOp.PrecisionFloat, UsdGeom.XformOp.PrecisionHalf]

        for prim_path in self.prim_paths:
            for precision in precision_types:
                await self.add_xform_ops_test(stage, prim_path, precision)

        for prim_path in self.prim_paths:
            for precision in precision_types:
                await self.add_xform_ops_test(stage, prim_path, precision, False)
                await self.remove_xform_ops_test(stage, prim_path)

    async def test_transform_add_remove_xformops_commands(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 wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()

        #ops_exists = []
        #for op_name in self.xform_op_names:
        #    ops_exists.append(self.include_op_code(xform, op_name))

        usd_context.get_selection().set_selected_prim_paths(["/World/Sphere"], True)
        await ui_test.human_delay(3)

        await self.add_xform_ops_test(stage, "/World/Sphere", UsdGeom.XformOp.PrecisionDouble, False)

        prim = stage.GetObjectAtPath("/World/Sphere")
        xform = UsdGeom.Xformable(prim)

        #disable all properties
        for op in self.xform_op_names:
            index = self.get_index_op_code(xform, op)
            omni.kit.commands.execute("RemoveXformOpCommand", op_order_attr_path="/World/Sphere.xformOpOrder", op_name=op, op_order_index=index)
            await ui_test.human_delay(3)
            self.assertTrue(not self.include_op_code(xform, op))

        #enable all properties
        for op in self.xform_op_names:
            omni.kit.commands.execute("EnableXformOpCommand", op_attr_path="/World/Sphere." + op)
            await ui_test.human_delay(3)
            self.assertTrue(self.include_op_code(xform, op))
            #undo
            omni.kit.commands.execute("Undo")
            await ui_test.human_delay(3)
            self.assertTrue(not self.include_op_code(xform, op))

        #disable all properties
        for op in self.xform_op_names:
            omni.kit.commands.execute("EnableXformOpCommand", op_attr_path="/World/Sphere." + op)
            await ui_test.human_delay(3)
            self.assertTrue(self.include_op_code(xform, op))

            index = self.get_index_op_code(xform, op)
            omni.kit.commands.execute("RemoveXformOpCommand", op_order_attr_path="/World/Sphere.xformOpOrder", op_name=op, op_order_index=index)
            await ui_test.human_delay(3)
            self.assertTrue(not self.include_op_code(xform, op))

            #undo
            omni.kit.commands.execute("Undo")
            await ui_test.human_delay(3)
            self.assertTrue(self.include_op_code(xform, op))

    async def test_chate_rotation_ops(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 wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)
        await ui_test.human_delay(3)

        command_rotation_ops = ["rotateX", "rotateY", "rotateZ", "rotateXYZ", "rotateXZY", "rotateYXZ", "rotateYZX", "rotateZXY", "rotateZYX", "rotateXYZ"]
        rotation_ops = ["rotateXZY", "rotateYXZ", "rotateYZX", "rotateZXY", "rotateZYX", "rotateXYZ"]
        #FIXME: if you fix rotateX, rotateY, rotateZ UI issue, uncomment below
        #rotation_ops = ["rotateX", "rotateY", "rotateZ", "rotateXYZ", "rotateXZY", "rotateYXZ", "rotateYZX", "rotateZXY", "rotateZYX", "rotateXYZ"]

        sphere_prim = stage.GetObjectAtPath("/World/Sphere")
        sphere_xform = UsdGeom.Xformable(sphere_prim)
        current = "rotateXYZ"
        for rotation_op in command_rotation_ops:
            omni.kit.commands.execute('ChangeRotationOp',
                src_op_attr_path=Sdf.Path('/World/Sphere.xformOp:'+current),
                op_name='xformOp:' + current,
                dst_op_attr_name='xformOp:' + rotation_op,
                is_inverse_op=False,
                auto_target_layer=True)

            self.assertTrue(self.include_op_code(sphere_xform, "xformOp:" + rotation_op))
            self.assertTrue(not self.include_op_code(sphere_xform, "xformOp:" + current))
            current = rotation_op

        #change rotation ops
        cube_prim = stage.GetObjectAtPath("/World/Cube")
        cube_xform = UsdGeom.Xformable(cube_prim)
        current = "rotateXYZ"
        for rotation_op in rotation_ops:
            omni.kit.commands.execute('ChangeRotationOp',
                src_op_attr_path=Sdf.Path('/World/Cube.xformOp:'+current),
                op_name='xformOp:' + current,
                dst_op_attr_name='xformOp:' + rotation_op,
                is_inverse_op=False,
                auto_target_layer=True)

            await ui_test.human_delay(3)
            self.assertTrue(self.include_op_code(cube_xform, "xformOp:" + rotation_op))
            self.assertTrue(not self.include_op_code(cube_xform, "xformOp:" + current))
            current = rotation_op

        #first set current value first
        test_value = Gf.Vec3d(30, 60, 90)
        attr_name ='xformOp:' + current
        attr = cube_prim.GetAttribute(attr_name)
        attr.Set(test_value)
        for rotation_op in rotation_ops:
            #create dest attribute first
            #and set it
            dest_attr_name = 'xformOp:' + rotation_op
            cube_prim.CreateAttribute(dest_attr_name, Sdf.ValueTypeNames.Double3,False).Set(Gf.Vec3d(0, 0, 180))
            omni.kit.commands.execute('ChangeRotationOp',
                src_op_attr_path=Sdf.Path('/World/Cube.xformOp:'+current),
                op_name='xformOp:' + current,
                dst_op_attr_name='xformOp:' + rotation_op,
                is_inverse_op=False,
                auto_target_layer=True)

            dest_attr = cube_prim.GetAttribute(dest_attr_name)
            # confirm the value is transferred
            diff = test_value - dest_attr.Get()
            self.assertTrue(diff.GetLength() < 0.1)

            # set new test value
            dest_attr.Set(test_value)

            current = rotation_op

    async def test_invalid_commands(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 wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()
        usd_context.get_selection().set_selected_prim_paths(["/World/Cube"], True)
        await ui_test.human_delay(3)

        #rotation_ops = ["rotateX", "rotateY", "rotateZ", "rotateXYZ", "rotateXZY", "rotateYXZ", "rotateYZX", "rotateZXY", "rotateZYX"]
        rotation_ops = ["rotateXZY", "rotateYXZ", "rotateYZX", "rotateZXY", "rotateZYX", "rotateXYZ"]

        #change rotation ops
        cube_prim = stage.GetObjectAtPath("/World/Cube")
        cube_xform = UsdGeom.Xformable(cube_prim)
        current = "rotateXYZ"

        #now test failing cases below
        for rotation_op in rotation_ops:
            omni.kit.commands.execute('ChangeRotationOp',
                src_op_attr_path=Sdf.Path('/World/Cube.xformOp:'+current),
                op_name='xformOp:' + current,
                dst_op_attr_name='xformOp:' + rotation_op,
                is_inverse_op=True,
                auto_target_layer=False)

        omni.kit.commands.execute('ChangeRotationOp',
            src_op_attr_path=Sdf.Path('/World/Cube.xformOp:rotateXYZ'),
            op_name='xformOp:rotateXYZ',
            dst_op_attr_name='/World/Cube.xformOp:rotateZYX',
            is_inverse_op=False,
            auto_target_layer=True)

        omni.kit.commands.execute('ChangeRotationOp',
            src_op_attr_path=Sdf.Path('/World/Cube.xformOp:rotateZYX'),
            op_name='xformOp:rotateZYX',
            dst_op_attr_name='xformOp:rotateXYZ',
            is_inverse_op=True,
            auto_target_layer=True)

        omni.kit.commands.execute("EnableXformOpCommand", op_attr_path="/World/Sphere.testAttribute")

        omni.kit.commands.execute('ChangeRotationOp',
            src_op_attr_path=Sdf.Path('/World/testAttribute.xformOp:rotateZYX'),
            op_name='xformOp:rotateG',
            dst_op_attr_name='xformOp:rotateL',
            is_inverse_op=False,
            auto_target_layer=False)

        payload = PrimSelectionPayload( weakref.ref(stage), [Sdf.Path("/World/Sphere")])
        omni.kit.commands.execute('AddXformOp',
            payload=payload,
            precision=UsdGeom.XformOp.TypeInvalid,
            rotation_order='XYZ',
            add_translate_op=True,
            add_rotateXYZ_op=True,
            add_orient_op=True,
            add_scale_op=True,
            add_transform_op=True,
            add_pivot_op=True)


    async def test_xform_op_utils(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 wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()
        usd_context.get_selection().set_selected_prim_paths(["/World/Sphere"], True)
        await ui_test.human_delay(3)

        self.assertTrue(get_op_type_name("Dummy") == None)

        # test _add_trs_op
        payload = PrimSelectionPayload( weakref.ref(stage), [Sdf.Path("/World/Sphere")])

        settings = carb.settings.get_settings()
        default_xform_op_precision = settings.get("/persistent/app/primCreation/DefaultXformOpPrecision")
        default_rotation_order = settings.get("/persistent/app/primCreation/DefaultRotationOrder")

        # set to double
        settings.set("/persistent/app/primCreation/DefaultXformOpPrecision", "Double")
        _add_trs_op(payload)
        omni.kit.commands.execute('Undo')

        settings.set("/persistent/app/primCreation/DefaultXformOpPrecision", "Float")
        _add_trs_op(payload)
        omni.kit.commands.execute('Undo')

        settings.set("/persistent/app/primCreation/DefaultXformOpPrecision", "Half")
        _add_trs_op(payload)
        omni.kit.commands.execute('Undo')

        #fail test
        settings.destroy_item("/persistent/app/primCreation/DefaultXformOpPrecision")
        settings.destroy_item("/persistent/app/primCreation/DefaultRotationOrder")
        _add_trs_op(payload)
        omni.kit.commands.execute('Undo')

        #recover default
        settings.set("/persistent/app/primCreation/DefaultXformOpPrecision", default_xform_op_precision)
        settings.set("/persistent/app/primCreation/DefaultRotationOrder", default_rotation_order)

        #get_op_precision
        self.assertTrue(get_op_precision(Sdf.ValueTypeNames.Quatf) == UsdGeom.XformOp.PrecisionFloat)
        self.assertTrue(get_op_precision(Sdf.ValueTypeNames.Matrix4d) == UsdGeom.XformOp.PrecisionDouble)
        self.assertTrue(get_op_precision(Sdf.ValueTypeNames.Half3) == UsdGeom.XformOp.PrecisionHalf)
        self.assertTrue(get_op_precision(Sdf.ValueTypeNames.Color3d) == UsdGeom.XformOp.PrecisionDouble)

        self.assertTrue(get_inverse_op_Name(RESET_XFORM_STACK, True) == RESET_XFORM_STACK)
        self.assertTrue(get_inverse_op_Name("Dummy", True) == "!invert!Dummy")
        self.assertTrue(get_inverse_op_Name("!invert!xformOp:Dummy", False) == "xformOp:Dummy")
        self.assertTrue(get_op_attr_name(RESET_XFORM_STACK) == None)