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

from __future__ import annotations

import asyncio
import math
import traceback
from enum import Enum, Flag, IntEnum, auto
from typing import Dict, List, Sequence, Set, Tuple, Union
import concurrent.futures

import carb
import carb.dictionary
import carb.events
import carb.profiler
import carb.settings
import omni.kit.app
from omni.kit.async_engine import run_coroutine
import omni.kit.commands
import omni.kit.undo
import omni.timeline
from omni.kit.manipulator.tool.snap import SnapProviderManager
from omni.kit.manipulator.tool.snap import settings_constants as snap_c
from omni.kit.manipulator.transform import AbstractTransformManipulatorModel, Operation
from omni.kit.manipulator.transform.gestures import (
    RotateChangedGesture,
    RotateDragGesturePayload,
    ScaleChangedGesture,
    ScaleDragGesturePayload,
    TransformDragGesturePayload,
    TranslateChangedGesture,
    TranslateDragGesturePayload,
)
from omni.kit.manipulator.transform.settings_constants import c
from omni.kit.manipulator.transform.settings_listener import OpSettingsListener, SnapSettingsListener
from omni.ui import scene as sc
from pxr import Gf, Sdf, Tf, Usd, UsdGeom, UsdUtils

from .settings_constants import Constants as prim_c
from .utils import *

# Settings needed for zero-gravity
TRANSFORM_GIZMO_CUSTOM_MANIPULATOR_ENABLED = "/app/transform/gizmoCustomManipulatorEnabled"
TRANSFORM_GIZMO_CUSTOM_MANIPULATOR_PRIMS = "/app/transform/gizmoCustomManipulatorPrims"
TRANSFORM_GIZMO_IS_USING = "/app/transform/gizmoIsUsing"
TRANSFORM_GIZMO_TRANSLATE_DELTA_XYZ = "/app/transform/gizmoTranslateDeltaXYZ"
TRANSFORM_GIZMO_PIVOT_WORLD_POSITION = "/app/transform/tempPivotWorldPosition"
TRANSFORM_GIZMO_ROTATE_DELTA_XYZW = "/app/transform/gizmoRotateDeltaXYZW"
TRANSFORM_GIZMO_SCALE_DELTA_XYZ = "/app/transform/gizmoScaleDeltaXYZ"
TRANSLATE_DELAY_FRAME_SETTING = "/exts/omni.kit.manipulator.prim/visual/delayFrame"


class OpFlag(Flag):
    TRANSLATE = auto()
    ROTATE = auto()
    SCALE = auto()


class Placement(Enum):
    LAST_PRIM_PIVOT = auto()
    SELECTION_CENTER = auto()
    BBOX_CENTER = auto()
    REF_PRIM = auto()
    BBOX_BASE = auto()


class ManipulationMode(IntEnum):
    PIVOT = 0  # transform around manipulator pivot
    UNIFORM = 1  # set same world transform from manipulator to all prims equally
    INDIVIDUAL = 2  # 2: (TODO) transform around each prim's own pivot respectively


class Viewport1WindowState:
    def __init__(self):
        self._focused_windows = None
        focused_windows = []
        try:
            # For some reason is_focused may return False, when a Window is definitely in fact is the focused window!
            # And there's no good solution to this when multiple Viewport-1 instances are open; so we just have to
            # operate on all Viewports for a given usd_context.
            import omni.kit.viewport_legacy as vp

            vpi = vp.acquire_viewport_interface()
            for instance in vpi.get_instance_list():
                window = vpi.get_viewport_window(instance)
                if not window:
                    continue
                focused_windows.append(window)
            if focused_windows:
                self._focused_windows = focused_windows
                for window in self._focused_windows:
                    # Disable the selection_rect, but enable_picking for snapping
                    window.disable_selection_rect(True)
                    # Schedule a picking request so if snap needs it later, it may arrive by the on_change event
                    window.request_picking()
        except Exception:
            pass

    def get_picked_world_pos(self):
        if self._focused_windows:
            # Try to reduce to the focused window now after, we've had some mouse-move input
            focused_windows = [window for window in self._focused_windows if window.is_focused()]
            if focused_windows:
                self._focused_windows = focused_windows
            for window in self._focused_windows:
                window.disable_selection_rect(True)
                # request picking FOR NEXT FRAME
                window.request_picking()
                # get PREVIOUSLY picked pos, it may be None the first frame but that's fine
                return window.get_picked_world_pos()
        return None

    def __del__(self):
        self.destroy()

    def destroy(self):
        self._focused_windows = None

    def get_usd_context_name(self):
        if self._focused_windows:
            return self._focused_windows[0].get_usd_context_name()
        else:
            return ""


class PrimTransformChangedGestureBase:
    def __init__(self, usd_context_name: str = "", viewport_api=None):
        self._settings = carb.settings.get_settings()
        self._usd_context_name = usd_context_name
        self._usd_context = omni.usd.get_context(self._usd_context_name)
        self._viewport_api = viewport_api  # VP2
        self._vp1_window_state = None
        self._stage_id = None

    def on_began(self, payload_type=TransformDragGesturePayload):
        self._viewport_on_began()

        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, payload_type):
            return

        model = self.sender.model
        if not model:
            return

        item = self.gesture_payload.changing_item
        self._current_editing_op = item.operation

        # NOTE! self._begin_xform has no scale. To get the full matrix, do self._begin_scale_mtx * self._begin_xform
        self._begin_xform = Gf.Matrix4d(*model.get_as_floats(model.get_item("no_scale_transform_manipulator")))

        manip_scale = Gf.Vec3d(*model.get_as_floats(model.get_item("scale_manipulator")))
        self._begin_scale_mtx = Gf.Matrix4d(1.0)
        self._begin_scale_mtx.SetScale(manip_scale)

        model.set_floats(model.get_item("viewport_fps"), [0.0])

        model.on_began(self.gesture_payload)

    def on_changed(self, payload_type=TransformDragGesturePayload):
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, payload_type):
            return

        model = self.sender.model
        if not model:
            return

        if self._viewport_api:
            fps = self._viewport_api.frame_info.get("fps")
            model.set_floats(model.get_item("viewport_fps"), [fps])

        model.on_changed(self.gesture_payload)

    def on_ended(self, payload_type=TransformDragGesturePayload):
        self._viewport_on_ended()
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, payload_type):
            return

        model = self.sender.model
        if not model:
            return

        item = self.gesture_payload.changing_item
        if item.operation != self._current_editing_op:
            return

        model.set_floats(model.get_item("viewport_fps"), [0.0])

        model.on_ended(self.gesture_payload)

        self._current_editing_op = None

    def on_canceled(self, payload_type=TransformDragGesturePayload):
        self._viewport_on_ended()
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, payload_type):
            return

        model = self.sender.model
        if not model:
            return

        item = self.gesture_payload.changing_item
        if item.operation != self._current_editing_op:
            return

        model.on_canceled(self.gesture_payload)

        self._current_editing_op = None

    def _publish_delta(self, operation: Operation, delta: List[float]):  # pragma: no cover
        if operation == Operation.TRANSLATE:
            self._settings.set_float_array(TRANSFORM_GIZMO_TRANSLATE_DELTA_XYZ, delta)
        elif operation == Operation.ROTATE:
            self._settings.set_float_array(TRANSFORM_GIZMO_ROTATE_DELTA_XYZW, delta)
        elif operation == Operation.SCALE:
            self._settings.set_float_array(TRANSFORM_GIZMO_SCALE_DELTA_XYZ, delta)

    def __set_viewport_manipulating(self, value: int):
        # Signal that user-manipulation has started for this stage
        if self._stage_id is None:
            self._stage_id = UsdUtils.StageCache.Get().GetId(self._usd_context.get_stage()).ToLongInt()
        key = f"/app/viewport/{self._stage_id}/manipulating"
        cur_value = self._settings.get(key) or 0
        self._settings.set(key, cur_value + value)

    def _viewport_on_began(self):
        self._viewport_on_ended()
        if self._viewport_api is None:
            self._vp1_window_state = Viewport1WindowState()
        self.__set_viewport_manipulating(1)

    def _viewport_on_ended(self):
        if self._vp1_window_state:
            self._vp1_window_state.destroy()
            self._vp1_window_state = None
        if self._stage_id:
            self.__set_viewport_manipulating(-1)
            self._stage_id = None


class PrimTranslateChangedGesture(TranslateChangedGesture, PrimTransformChangedGestureBase):
    def __init__(self, snap_manager: SnapProviderManager, **kwargs):
        PrimTransformChangedGestureBase.__init__(self, **kwargs)
        TranslateChangedGesture.__init__(self)
        self._accumulated_translate = Gf.Vec3d(0)
        self._snap_manager = snap_manager

    def on_began(self):
        PrimTransformChangedGestureBase.on_began(self, TranslateDragGesturePayload)
        self._accumulated_translate = Gf.Vec3d(0)

        model = self._get_model(TranslateDragGesturePayload)
        if model and self._can_snap(model):
            # TODO No need for gesture=self when VP1 has viewport_api
            self._snap_manager.on_began(model.consolidated_xformable_prim_data_curr.keys(), gesture=self)

        if model:
            model.set_floats(model.get_item("translate_delta"), [0, 0, 0])

    def on_ended(self):
        PrimTransformChangedGestureBase.on_ended(self, TranslateDragGesturePayload)

        model = self._get_model(TranslateDragGesturePayload)
        if model and self._can_snap(model):
            self._snap_manager.on_ended()

    def on_canceled(self):
        PrimTransformChangedGestureBase.on_canceled(self, TranslateDragGesturePayload)

        model = self._get_model(TranslateDragGesturePayload)
        if model and self._can_snap(model):
            self._snap_manager.on_ended()

    @carb.profiler.profile
    def on_changed(self):
        PrimTransformChangedGestureBase.on_changed(self, TranslateDragGesturePayload)

        model = self._get_model(TranslateDragGesturePayload)
        if not model:
            return

        manip_xform = Gf.Matrix4d(*model.get_as_floats(model.get_item("no_scale_transform_manipulator")))
        new_manip_xform = Gf.Matrix4d(manip_xform)
        rotation_mtx = manip_xform.ExtractRotationMatrix()
        rotation_mtx.Orthonormalize()

        translate_delta = self.gesture_payload.moved_delta
        translate = self.gesture_payload.moved
        axis = self.gesture_payload.axis

        # slow Gf.Vec3d(*translate_delta)
        translate_delta = Gf.Vec3d(translate_delta[0], translate_delta[1], translate_delta[2])
        if model.op_settings_listener.translation_mode == c.TRANSFORM_MODE_LOCAL:
            translate_delta = translate_delta * rotation_mtx

        self._accumulated_translate += translate_delta

        def apply_position(snap_world_pos=None, snap_world_orient=None, keep_spacing: bool = True):
            nonlocal new_manip_xform
            if self.state != sc.GestureState.CHANGED:
                return

            # only set translate if no snap or only snap to position
            item_name = "translate"

            if snap_world_pos and (
                math.isfinite(snap_world_pos[0])
                and math.isfinite(snap_world_pos[1])
                and math.isfinite(snap_world_pos[2])
            ):
                if snap_world_orient is None:
                    new_manip_xform.SetTranslateOnly(Gf.Vec3d(snap_world_pos[0], snap_world_pos[1], snap_world_pos[2]))
                    new_manip_xform = self._begin_scale_mtx * new_manip_xform
                else:
                    new_manip_xform.SetTranslateOnly(Gf.Vec3d(snap_world_pos[0], snap_world_pos[1], snap_world_pos[2]))
                    new_manip_xform.SetRotateOnly(snap_world_orient)

                    # set transform if snap both position and orientation
                    item_name = "no_scale_transform_manipulator"
            else:
                new_manip_xform.SetTranslateOnly(self._begin_xform.ExtractTranslation() + self._accumulated_translate)
                new_manip_xform = self._begin_scale_mtx * new_manip_xform

            model.set_floats(model.get_item("translate_delta"), translate_delta)

            if model.custom_manipulator_enabled:
                self._publish_delta(Operation.TRANSLATE, translate_delta)

            if keep_spacing is False:
                mode_item = model.get_item("manipulator_mode")
                prev_mode = model.get_as_ints(mode_item)
                model.set_ints(mode_item, [int(ManipulationMode.UNIFORM)])

            model.set_floats(model.get_item(item_name), flatten(new_manip_xform))

            if keep_spacing is False:
                model.set_ints(mode_item, prev_mode)

        # only do snap to surface if drag the center point
        if (
            model.snap_settings_listener.snap_enabled
            and model.snap_settings_listener.snap_provider
            and axis == [1, 1, 1]
        ):
            ndc_location = None
            if self._viewport_api:
                # No mouse location is available, have to convert back to NDC space
                ndc_location = self.sender.transform_space(
                    sc.Space.WORLD, sc.Space.NDC, self.gesture_payload.ray_closest_point
                )

            if self._snap_manager.get_snap_pos(
                new_manip_xform,
                ndc_location,
                self.sender.scene_view,
                lambda **kwargs: apply_position(
                    kwargs.get("position", None), kwargs.get("orient", None), kwargs.get("keep_spacing", True)
                ),
            ):
                return

        apply_position()

    def _get_model(self, payload_type) -> PrimTransformModel:
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, payload_type):
            return None

        return self.sender.model

    def _can_snap(self, model: PrimTransformModel):
        axis = self.gesture_payload.axis
        if (
            model.snap_settings_listener.snap_enabled
            and model.snap_settings_listener.snap_provider
            and axis == [1, 1, 1]
        ):
            return True

        return False


class PrimRotateChangedGesture(RotateChangedGesture, PrimTransformChangedGestureBase):
    def __init__(self, **kwargs):
        PrimTransformChangedGestureBase.__init__(self, **kwargs)
        RotateChangedGesture.__init__(self)

    def on_began(self):
        PrimTransformChangedGestureBase.on_began(self, RotateDragGesturePayload)

        model = self.sender.model
        if model:
            model.set_floats(model.get_item("rotate_delta"), [0, 0, 0, 0])

    def on_ended(self):
        PrimTransformChangedGestureBase.on_ended(self, RotateDragGesturePayload)

    def on_canceled(self):
        PrimTransformChangedGestureBase.on_canceled(self, RotateDragGesturePayload)

    @carb.profiler.profile
    def on_changed(self):
        PrimTransformChangedGestureBase.on_changed(self, RotateDragGesturePayload)

        if (
            not self.gesture_payload
            or not self.sender
            or not isinstance(self.gesture_payload, RotateDragGesturePayload)
        ):
            return

        model = self.sender.model
        if not model:
            return

        axis = self.gesture_payload.axis
        angle = self.gesture_payload.angle
        angle_delta = self.gesture_payload.angle_delta
        screen_space = self.gesture_payload.screen_space
        free_rotation = self.gesture_payload.free_rotation

        axis = Gf.Vec3d(*axis[:3])
        rotate = Gf.Rotation(axis, angle)
        delta_axis = Gf.Vec4d(*axis, 0.0)

        rot_matrix = Gf.Matrix4d(1)
        rot_matrix.SetRotate(rotate)

        if free_rotation:
            rotate = Gf.Rotation(axis, angle_delta)

            rot_matrix = Gf.Matrix4d(1)
            rot_matrix.SetRotate(rotate)

            xform = Gf.Matrix4d(*model.get_as_floats(model.get_item("no_scale_transform_manipulator")))
            full_xform = self._begin_scale_mtx * xform
            translate = full_xform.ExtractTranslation()
            no_translate_mtx = Gf.Matrix4d(full_xform)
            no_translate_mtx.SetTranslateOnly(Gf.Vec3d(0))
            no_translate_mtx = no_translate_mtx * rot_matrix
            new_transform_matrix = no_translate_mtx.SetTranslateOnly(translate)
            delta_axis = no_translate_mtx * delta_axis
        elif model.op_settings_listener.rotation_mode == c.TRANSFORM_MODE_GLOBAL or screen_space:
            begin_full_xform = self._begin_scale_mtx * self._begin_xform
            translate = begin_full_xform.ExtractTranslation()
            no_translate_mtx = Gf.Matrix4d(begin_full_xform)
            no_translate_mtx.SetTranslateOnly(Gf.Vec3d(0))
            no_translate_mtx = no_translate_mtx * rot_matrix
            new_transform_matrix = no_translate_mtx.SetTranslateOnly(translate)
            delta_axis = no_translate_mtx * delta_axis
        else:
            self._begin_xform = Gf.Matrix4d(*model.get_as_floats(model.get_item("no_scale_transform_manipulator")))
            new_transform_matrix = self._begin_scale_mtx * rot_matrix * self._begin_xform

        delta_axis.Normalize()
        delta_rotate = Gf.Rotation(delta_axis[:3], angle_delta)
        quat = delta_rotate.GetQuaternion()
        real = quat.GetReal()
        imaginary = quat.GetImaginary()
        rd = [imaginary[0], imaginary[1], imaginary[2], real]
        model.set_floats(model.get_item("rotate_delta"), rd)

        if model.custom_manipulator_enabled:
            self._publish_delta(Operation.ROTATE, rd)

        model.set_floats(model.get_item("rotate"), flatten(new_transform_matrix))


class PrimScaleChangedGesture(ScaleChangedGesture, PrimTransformChangedGestureBase):
    def __init__(self, **kwargs):
        PrimTransformChangedGestureBase.__init__(self, **kwargs)
        ScaleChangedGesture.__init__(self)

    def on_began(self):
        PrimTransformChangedGestureBase.on_began(self, ScaleDragGesturePayload)

        model = self.sender.model
        if model:
            model.set_floats(model.get_item("scale_delta"), [0, 0, 0])

    def on_ended(self):
        PrimTransformChangedGestureBase.on_ended(self, ScaleDragGesturePayload)

    def on_canceled(self):
        PrimTransformChangedGestureBase.on_canceled(self, ScaleDragGesturePayload)

    @carb.profiler.profile
    def on_changed(self):
        PrimTransformChangedGestureBase.on_changed(self, ScaleDragGesturePayload)

        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, ScaleDragGesturePayload):
            return

        model = self.sender.model
        if not model:
            return

        axis = self.gesture_payload.axis
        scale = self.gesture_payload.scale

        axis = Gf.Vec3d(*axis[:3])

        scale_delta = scale * axis
        scale_vec = Gf.Vec3d()
        for i in range(3):
            scale_vec[i] = scale_delta[i] if scale_delta[i] else 1

        scale_matrix = Gf.Matrix4d(1.0)
        scale_matrix.SetScale(scale_vec)
        scale_matrix *= self._begin_scale_mtx
        new_transform_matrix = scale_matrix * self._begin_xform

        s = Gf.Vec3d(*model.get_as_floats(model.get_item("scale_manipulator")))
        sd = [s_n / s_o for s_n, s_o in zip([scale_matrix[0][0], scale_matrix[1][1], scale_matrix[2][2]], s)]
        model.set_floats(model.get_item("scale_delta"), sd)

        if model.custom_manipulator_enabled:
            self._publish_delta(Operation.SCALE, sd)

        model.set_floats(model.get_item("scale"), flatten(new_transform_matrix))


class PrimTransformModel(AbstractTransformManipulatorModel):
    def __init__(self, usd_context_name: str = "", viewport_api=None):
        super().__init__()
        self._transform = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]  # visual transform of manipulator
        self._no_scale_transform_manipulator = Gf.Matrix4d(1.0)  # transform of manipulator without scale
        self._scale_manipulator = Gf.Vec3d(1.0)  # scale of manipulator
        self._settings = carb.settings.get_settings()
        self._dict = carb.dictionary.get_dictionary()
        self._timeline = omni.timeline.get_timeline_interface()
        self._app = omni.kit.app.get_app()
        self._usd_context_name = usd_context_name
        self._usd_context = omni.usd.get_context(usd_context_name)

        self._stage: Usd.Stage = None
        self._enabled_hosting_widget_count: int = 0
        self._stage_listener = None
        self._xformable_prim_paths: List[Sdf.Path] = []
        self._xformable_prim_paths_sorted: List[Sdf.Path] = []
        self._xformable_prim_paths_set: Set[Sdf.Path] = set()
        self._xformable_prim_paths_prefix_set: Set[Sdf.Path] = set()
        self._consolidated_xformable_prim_paths: List[Sdf.Path] = []
        self._pivot_prim: Usd.Prim = None
        self._update_prim_xform_from_prim_task_or_future: Union[asyncio.Task, concurrent.futures.Future, None] = None
        self._current_editing_op: Operation = None
        self._ignore_xform_data_change = False
        self._timeline_sub = None
        self._pending_changed_paths: Dict[Sdf.Path, bool] = {}
        self._pending_changed_paths_for_xform_data: Set[Sdf.Path] = set()
        self._mode = ManipulationMode.PIVOT
        self._viewport_fps: float = 0.0  # needed this as heuristic for delaying the visual update of manipulator to match rendering
        self._viewport_api = viewport_api
        self._delay_dirty_tasks_or_futures: Dict[int, Union[asyncio.Task, concurrent.futures.Future]] = {}

        self._no_scale_transform_manipulator_item = sc.AbstractManipulatorItem()
        self._items["no_scale_transform_manipulator"] = self._no_scale_transform_manipulator_item

        self._scale_manipulator_item = sc.AbstractManipulatorItem()
        self._items["scale_manipulator"] = self._scale_manipulator_item

        self._transform_manipulator_item = sc.AbstractManipulatorItem()
        self._items["transform_manipulator"] = self._transform_manipulator_item

        self._manipulator_mode_item = sc.AbstractManipulatorItem()
        self._items["manipulator_mode"] = self._manipulator_mode_item

        self._viewport_fps_item = sc.AbstractManipulatorItem()
        self._items["viewport_fps"] = self._viewport_fps_item

        self._set_default_settings()

        # subscribe to snap events
        self._snap_settings_listener = SnapSettingsListener(
            enabled_setting_path=None,
            move_x_setting_path=snap_c.SNAP_TRANSLATE_SETTING_PATH,
            move_y_setting_path=snap_c.SNAP_TRANSLATE_SETTING_PATH,
            move_z_setting_path=snap_c.SNAP_TRANSLATE_SETTING_PATH,
            rotate_setting_path=snap_c.SNAP_ROTATE_SETTING_PATH,
            scale_setting_path=snap_c.SNAP_SCALE_SETTING_PATH,
            provider_setting_path=snap_c.SNAP_PROVIDER_NAME_SETTING_PATH,
        )

        # subscribe to operation/mode events
        self._op_settings_listener = OpSettingsListener()
        self._op_settings_listener_sub = self._op_settings_listener.subscribe_listener(self._on_op_listener_changed)

        self._placement_sub = self._settings.subscribe_to_node_change_events(
            prim_c.MANIPULATOR_PLACEMENT_SETTING, self._on_placement_setting_changed
        )
        self._placement = Placement.LAST_PRIM_PIVOT
        placement = self._settings.get(prim_c.MANIPULATOR_PLACEMENT_SETTING)
        self._update_placement(placement)

        # cache unique setting path for selection pivot position
        # vp1 default
        self._selection_pivot_position_path = TRANSFORM_GIZMO_PIVOT_WORLD_POSITION + "/Viewport"
        # vp2
        if self._viewport_api:
            self._selection_pivot_position_path = (
                TRANSFORM_GIZMO_PIVOT_WORLD_POSITION + "/" + self._viewport_api.id.split("/")[0]
            )
        # update setting with pivot placement position on init
        self._update_temp_pivot_world_position()

        def subscribe_to_value_and_get_current(setting_val_name: str, setting_path: str):
            sub = self._settings.subscribe_to_node_change_events(
                setting_path, lambda item, type: setattr(self, setting_val_name, self._dict.get(item))
            )
            setattr(self, setting_val_name, self._settings.get(setting_path))
            return sub

        # subscribe to zero-gravity settings
        self._custom_manipulator_enabled_sub = subscribe_to_value_and_get_current(
            "_custom_manipulator_enabled", TRANSFORM_GIZMO_CUSTOM_MANIPULATOR_ENABLED
        )

        # subscribe to USD related events
        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop(
            self._on_stage_event, name="PrimTransformModel stage event"
        )

        if self._usd_context.get_stage_state() == omni.usd.StageState.OPENED:
            self._on_stage_opened()

        self._warning_notification = None
        self._selected_instance_proxy_paths = set()

    def __del__(self):
        self.destroy()

    def destroy(self):
        if self._warning_notification:
            self._warning_notification.dismiss()
            self._warning_notification = None

        self._op_settings_listener_sub = None
        if self._op_settings_listener:
            self._op_settings_listener.destroy()
            self._op_settings_listener = None

        self._stage_event_sub = None
        if self._stage_listener:
            self._stage_listener.Revoke()
            self._stage_listener = None

        if self._usd_context.get_stage_state() == omni.usd.StageState.OPENED:
            self._on_stage_closing()

        self._timeline_sub = None

        if self._snap_settings_listener:
            self._snap_settings_listener.destroy()
            self._snap_settings_listener = None

        if self._custom_manipulator_enabled_sub:
            self._settings.unsubscribe_to_change_events(self._custom_manipulator_enabled_sub)
            self._custom_manipulator_enabled_sub = None

        if self._placement_sub:
            self._settings.unsubscribe_to_change_events(self._placement_sub)
            self._placement_sub = None

        for task_or_future in self._delay_dirty_tasks_or_futures.values():
            task_or_future.cancel()
        self._delay_dirty_tasks_or_futures.clear()

    def set_pivot_prim_path(self, path: Sdf.Path) -> bool:
        if path not in self._xformable_prim_paths_set:
            carb.log_warn(f"Cannot set pivot prim path to {path}")
            return False

        if not self._pivot_prim or self._pivot_prim.GetPath() != path:
            self._pivot_prim = self._stage.GetPrimAtPath(path)
        else:
            return False

        if self._update_transform_from_prims():
            self._item_changed(self._transform_item)

        return True

    def get_pivot_prim_path(self) -> Sdf.Path:
        if self._pivot_prim:
            return self._pivot_prim.GetPath()

        return None

    def on_began(self, payload):
        item = payload.changing_item
        self._current_editing_op = item.operation

        # All selected xformable prims' transforms. Store this for when `Keep Spacing` option is off during snapping,
        # because it can modify parent's and child're transform differently.
        self._all_xformable_prim_data_prev: Dict[Sdf.Path, Tuple[Gf.Vec3d, Gf.Vec3d, Gf.Vec3i, Gf.Vec3d, Gf.Vec3d]] = {}

        # consolidated xformable prims' transforms. If parent is in the dict, child will be excluded.
        self._consolidated_xformable_prim_data_prev: Dict[Sdf.Path, Tuple[Gf.Vec3d, Gf.Vec3d, Gf.Vec3i, Gf.Vec3d]] = {}

        for path in self._xformable_prim_paths:
            prim = self.stage.GetPrimAtPath(path)
            xform_tuple = omni.usd.get_local_transform_SRT(prim, self._get_current_time_code())
            pivot = get_local_transform_pivot_inv(prim, self._get_current_time_code()).GetInverse()
            self._all_xformable_prim_data_prev[path] = xform_tuple + (pivot,)

            if path in self._consolidated_xformable_prim_paths:
                self._consolidated_xformable_prim_data_prev[path] = xform_tuple + (pivot,)

        self.all_xformable_prim_data_curr = self._all_xformable_prim_data_prev.copy()
        self.consolidated_xformable_prim_data_curr = self._consolidated_xformable_prim_data_prev.copy()

        self._pending_changed_paths_for_xform_data.clear()

        if self.custom_manipulator_enabled:
            self._settings.set(TRANSFORM_GIZMO_IS_USING, True)

    def on_changed(self, payload):
        # Always re-fetch the changed transform on_changed. Although it might not be manipulated by manipulator,
        # prim's transform can be modified by other runtime (e.g. omnigraph), and we always want the latest.
        self._update_xform_data_from_dirty_paths()

    def on_ended(self, payload):
        # Always re-fetch the changed transform on_changed. Although it might not be manipulated by manipulator,
        # prim's transform can be modified by other runtime (e.g. omnigraph), and we always want the latest.
        self._update_xform_data_from_dirty_paths()

        carb.profiler.begin(1, "PrimTransformChangedGestureBase.TransformPrimSRT.all")
        paths = []
        new_translations = []
        new_rotation_eulers = []
        new_rotation_orders = []
        new_scales = []
        old_translations = []
        old_rotation_eulers = []
        old_rotation_orders = []
        old_scales = []

        for path, (s, r, ro, t, pivot) in self.all_xformable_prim_data_curr.items():
            # Data didn't change
            if self._all_xformable_prim_data_prev[path] == self.all_xformable_prim_data_curr[path]:
                # carb.log_info(f"Skip {path}")
                continue

            (old_s, old_r, old_ro, old_t, old_pivot) = self._all_xformable_prim_data_prev[path]

            paths.append(path.pathString)
            new_translations += [t[0], t[1], t[2]]
            new_rotation_eulers += [r[0], r[1], r[2]]
            new_rotation_orders += [ro[0], ro[1], ro[2]]
            new_scales += [s[0], s[1], s[2]]
            old_translations += [old_t[0], old_t[1], old_t[2]]
            old_rotation_eulers += [old_r[0], old_r[1], old_r[2]]
            old_rotation_orders += [old_ro[0], old_ro[1], old_ro[2]]
            old_scales += [old_s[0], old_s[1], old_s[2]]

        self._ignore_xform_data_change = True
        self._on_ended_transform(
            paths,
            new_translations,
            new_rotation_eulers,
            new_rotation_orders,
            new_scales,
            old_translations,
            old_rotation_eulers,
            old_rotation_orders,
            old_scales,
        )
        self._ignore_xform_data_change = False

        carb.profiler.end(1)

        if self.custom_manipulator_enabled:
            self._settings.set(TRANSFORM_GIZMO_IS_USING, False)

        # if the manipulator was locked to orientation or translation,
        # refresh it on_ended so the transform is up to date
        mode = self._get_transform_mode_for_current_op()
        if self._should_keep_manipulator_orientation_unchanged(
            mode
        ) or self._should_keep_manipulator_translation_unchanged(mode):
            # set editing op to None AFTER _should_keep_manipulator_*_unchanged but
            # BEFORE self._update_transform_from_prims
            self._current_editing_op = None
            if self._update_transform_from_prims():
                self._item_changed(self._transform_item)

        self._current_editing_op = None

    def _on_ended_transform(
        self,
        paths: List[str],
        new_translations: List[float],
        new_rotation_eulers: List[float],
        new_rotation_orders: List[int],
        new_scales: List[float],
        old_translations: List[float],
        old_rotation_eulers: List[float],
        old_rotation_orders: List[int],
        old_scales: List[float],
    ):
        """Function ran by on_ended(). Can be overridden to change the behavior. Do not remove this function: can be
        used outside of this code"""
        self._alert_if_selection_has_instance_proxies()
        omni.kit.commands.execute(
            "TransformMultiPrimsSRTCpp",
            count=len(paths),
            paths=paths,
            new_translations=new_translations,
            new_rotation_eulers=new_rotation_eulers,
            new_rotation_orders=new_rotation_orders,
            new_scales=new_scales,
            old_translations=old_translations,
            old_rotation_eulers=old_rotation_eulers,
            old_rotation_orders=old_rotation_orders,
            old_scales=old_scales,
            usd_context_name=self._usd_context_name,
            time_code=self._get_current_time_code().GetValue(),
        )

    def on_canceled(self, payload):
        if self.custom_manipulator_enabled:
            self._settings.set(TRANSFORM_GIZMO_IS_USING, False)

        self._current_editing_op = None

    def widget_enabled(self):
        self._enabled_hosting_widget_count += 1

        # just changed from no active widget to 1 active widget
        if self._enabled_hosting_widget_count == 1:
            # listener only needs to be activated if manipulator is visible
            if self._consolidated_xformable_prim_paths:
                assert self._stage_listener is None
                self._stage_listener = Tf.Notice.Register(
                    Usd.Notice.ObjectsChanged, self._on_objects_changed, self._stage
                )
                carb.log_info("Tf.Notice.Register in PrimTransformModel")

    def _clear_temp_pivot_position_setting(self):
        if self._settings.get(self._selection_pivot_position_path):
            self._settings.destroy_item(self._selection_pivot_position_path)

    def widget_disabled(self):
        self._enabled_hosting_widget_count -= 1
        self._clear_temp_pivot_position_setting()
        assert self._enabled_hosting_widget_count >= 0
        if self._enabled_hosting_widget_count < 0:
            carb.log_error(f"manipulator enabled widget tracker out of sync: {self._enabled_hosting_widget_count}")
            self._enabled_hosting_widget_count = 0

        # If no hosting manipulator is active, revoke the listener since there's no need to sync Transform
        if self._enabled_hosting_widget_count == 0:
            if self._stage_listener:
                self._stage_listener.Revoke()
                self._stage_listener = None
                carb.log_info("Tf.Notice.Revoke in PrimTransformModel")

            for task_or_future in self._delay_dirty_tasks_or_futures.values():
                task_or_future.cancel()
            self._delay_dirty_tasks_or_futures.clear()

    def set_floats(self, item: sc.AbstractManipulatorItem, value: Sequence[float]):
        if item == self._viewport_fps_item:
            self._viewport_fps = value[0]
            return

        flag = None
        if issubclass(type(item), AbstractTransformManipulatorModel.OperationItem):
            if (
                item.operation == Operation.TRANSLATE_DELTA
                or item.operation == Operation.ROTATE_DELTA
                or item.operation == Operation.SCALE_DELTA
            ):
                return
            if item.operation == Operation.TRANSLATE:
                flag = OpFlag.TRANSLATE
            elif item.operation == Operation.ROTATE:
                flag = OpFlag.ROTATE
            elif item.operation == Operation.SCALE:
                flag = OpFlag.SCALE
            transform = Gf.Matrix4d(*value)
        elif item == self._transform_manipulator_item:
            flag = OpFlag.TRANSLATE | OpFlag.ROTATE | OpFlag.SCALE
            transform = Gf.Matrix4d(*value)
        elif item == self._no_scale_transform_manipulator_item:
            flag = OpFlag.TRANSLATE | OpFlag.ROTATE
            old_manipulator_scale_mtx = Gf.Matrix4d(1.0)
            old_manipulator_scale_mtx.SetScale(self._scale_manipulator)
            transform = old_manipulator_scale_mtx * Gf.Matrix4d(*value)

        if flag is not None:
            if self._mode == ManipulationMode.PIVOT:
                self._transform_selected_prims(
                    transform, self._no_scale_transform_manipulator, self._scale_manipulator, flag
                )
            elif self._mode == ManipulationMode.UNIFORM:
                self._transform_all_selected_prims_to_manipulator_pivot(transform, flag)
        else:
            carb.log_warn(f"Unsupported item {item}")

    def set_ints(self, item: sc.AbstractManipulatorItem, value: Sequence[int]):
        if item == self._manipulator_mode_item:
            try:
                self._mode = ManipulationMode(value[0])
            except:
                carb.log_error(traceback.format_exc())
        else:
            carb.log_warn(f"unsupported item {item}")
            return None

    def get_as_floats(self, item: sc.AbstractManipulatorItem):
        if item == self._transform_item:
            return self._transform
        elif item == self._no_scale_transform_manipulator_item:
            return flatten(self._no_scale_transform_manipulator)
        elif item == self._scale_manipulator_item:
            return [self._scale_manipulator[0], self._scale_manipulator[1], self._scale_manipulator[2]]
        elif item == self._transform_manipulator_item:
            scale_mtx = Gf.Matrix4d(1)
            scale_mtx.SetScale(self._scale_manipulator)
            return flatten(scale_mtx * self._no_scale_transform_manipulator)
        else:
            carb.log_warn(f"unsupported item {item}")
            return None

    def get_as_ints(self, item: sc.AbstractManipulatorItem):
        if item == self._manipulator_mode_item:
            return [int(self._mode)]
        else:
            carb.log_warn(f"unsupported item {item}")
            return None

    def get_operation(self) -> Operation:
        if self._op_settings_listener.selected_op == c.TRANSFORM_OP_MOVE:
            return Operation.TRANSLATE
        elif self._op_settings_listener.selected_op == c.TRANSFORM_OP_ROTATE:
            return Operation.ROTATE
        elif self._op_settings_listener.selected_op == c.TRANSFORM_OP_SCALE:
            return Operation.SCALE
        return Operation.NONE

    def get_snap(self, item: AbstractTransformManipulatorModel.OperationItem):
        if not self._snap_settings_listener.snap_enabled:
            return None

        if item.operation == Operation.TRANSLATE:
            if self._snap_settings_listener.snap_provider:
                return None

            return (
                self._snap_settings_listener.snap_move_x,
                self._snap_settings_listener.snap_move_y,
                self._snap_settings_listener.snap_move_z,
            )
        elif item.operation == Operation.ROTATE:
            return self._snap_settings_listener.snap_rotate
        elif item.operation == Operation.SCALE:
            return self._snap_settings_listener.snap_scale

        return None

    @property
    def stage(self):
        return self._stage

    @property
    def custom_manipulator_enabled(self):
        return self._custom_manipulator_enabled

    @property
    def snap_settings_listener(self):
        return self._snap_settings_listener

    @property
    def op_settings_listener(self):
        return self._op_settings_listener

    @property
    def usd_context(self) -> omni.usd.UsdContext:
        return self._usd_context

    @property
    def xformable_prim_paths(self) -> List[Sdf.Path]:
        return self._xformable_prim_paths

    @carb.profiler.profile
    def _update_xform_data_from_dirty_paths(self):
        for p in self._pending_changed_paths_for_xform_data:
            prim_path = p.GetPrimPath()
            if prim_path in self.all_xformable_prim_data_curr and self._path_may_affect_transform(p):
                prim = self.stage.GetPrimAtPath(prim_path)
                xform_tuple = omni.usd.get_local_transform_SRT(prim, self._get_current_time_code())
                pivot = get_local_transform_pivot_inv(prim, self._get_current_time_code()).GetInverse()
                self.all_xformable_prim_data_curr[prim_path] = xform_tuple + (pivot,)

                if prim_path in self._consolidated_xformable_prim_paths:
                    self.consolidated_xformable_prim_data_curr[prim_path] = xform_tuple + (pivot,)

        self._pending_changed_paths_for_xform_data.clear()

    @carb.profiler.profile
    def _transform_selected_prims(
        self,
        new_manipulator_transform: Gf.Matrix4d,
        old_manipulator_transform_no_scale: Gf.Matrix4d,
        old_manipulator_scale: Gf.Vec3d,
        dirty_ops: OpFlag,
    ):
        carb.profiler.begin(1, "omni.kit.manipulator.prim.model._transform_selected_prims.prepare_data")
        paths = []
        new_translations = []
        new_rotation_eulers = []
        new_rotation_orders = []
        new_scales = []

        self._xform_cache.Clear()

        # any op may trigger a translation change if multi-manipulating
        should_update_translate = dirty_ops & OpFlag.TRANSLATE or len(
            self._xformable_prim_paths) > 1 or self._placement != Placement.LAST_PRIM_PIVOT
        should_update_rotate = dirty_ops & OpFlag.ROTATE
        should_update_scale = dirty_ops & OpFlag.SCALE

        old_manipulator_scale_mtx = Gf.Matrix4d(1.0)
        old_manipulator_scale_mtx.SetScale(old_manipulator_scale)
        old_manipulator_transform_inv = (old_manipulator_scale_mtx * old_manipulator_transform_no_scale).GetInverse()

        for path in self._consolidated_xformable_prim_paths:
            if self._custom_manipulator_enabled and self._should_skip_custom_manipulator_path(path.pathString):
                continue

            selected_prim = self._stage.GetPrimAtPath(path)
            # We check whether path is in consolidated_xformable_prim_data_curr because it may have not made it the dictionary if an error occured
            if not selected_prim or path not in self.consolidated_xformable_prim_data_curr:
                continue

            (s, r, ro, t, selected_pivot) = self.consolidated_xformable_prim_data_curr[path]
            selected_pivot_inv = selected_pivot.GetInverse()

            selected_local_to_world_mtx = self._xform_cache.GetLocalToWorldTransform(selected_prim)
            selected_parent_to_world_mtx = self._xform_cache.GetParentToWorldTransform(selected_prim)

            # Transform the prim from world space to pivot's space
            # Then apply the new pivotPrim-to-world-space transform matrix
            selected_local_to_world_pivot_mtx = (
                selected_pivot * selected_local_to_world_mtx * old_manipulator_transform_inv * new_manipulator_transform
            )

            world_to_parent_mtx = selected_parent_to_world_mtx.GetInverse()
            selected_local_mtx_new = selected_local_to_world_pivot_mtx * world_to_parent_mtx * selected_pivot_inv

            if should_update_translate:
                translation = selected_local_mtx_new.ExtractTranslation()

            if should_update_rotate:
                # Construct the new rotation from old scale and translation.
                # Don't use Factor because it won't be able to tell if scale is positive or negative and can result in flipped rotation
                old_s_mtx = Gf.Matrix4d(1.0)
                old_s_mtx.SetScale(Gf.Vec3d(s))
                old_t_mtx = Gf.Matrix4d(1.0)
                old_t_mtx.SetTranslate(Gf.Vec3d(t))
                rot_new = (old_s_mtx.GetInverse() * selected_local_mtx_new * old_t_mtx.GetInverse()).ExtractRotation()

                axes = [Gf.Vec3d.XAxis(), Gf.Vec3d.YAxis(), Gf.Vec3d.ZAxis()]
                decomp_rot = rot_new.Decompose(axes[ro[2]], axes[ro[1]], axes[ro[0]])
                index_order = Gf.Vec3i()
                for i in range(3):
                    index_order[ro[i]] = 2 - i

                rotation = Gf.Vec3d(decomp_rot[index_order[0]], decomp_rot[index_order[1]], decomp_rot[index_order[2]])
                rotation = find_best_euler_angles(Gf.Vec3d(r), rotation, ro)

            if should_update_scale:
                # Construct the new scale from old rotation and translation.
                # Don't use Factor because it won't be able to tell if scale is positive or negative and can result in flipped rotation
                old_rt_mtx = compose_transform_ops_to_matrix(t, r, ro, Gf.Vec3d(1))
                new_s_mtx = selected_local_mtx_new * old_rt_mtx.GetInverse()
                scale = Gf.Vec3d(new_s_mtx[0][0], new_s_mtx[1][1], new_s_mtx[2][2])

            translation = translation if should_update_translate else t
            rotation = rotation if should_update_rotate else r
            scale = scale if should_update_scale else s

            paths.append(path.pathString)
            new_translations += [translation[0], translation[1], translation[2]]
            new_rotation_eulers += [rotation[0], rotation[1], rotation[2]]
            new_rotation_orders += [ro[0], ro[1], ro[2]]
            new_scales += [scale[0], scale[1], scale[2]]

            xform_tuple = (scale, rotation, ro, translation, selected_pivot)
            self.consolidated_xformable_prim_data_curr[path] = xform_tuple
            self.all_xformable_prim_data_curr[path] = xform_tuple
        carb.profiler.end(1)

        self._ignore_xform_data_change = True
        self._do_transform_selected_prims(paths, new_translations, new_rotation_eulers, new_rotation_orders, new_scales)
        self._ignore_xform_data_change = False

    def _do_transform_selected_prims(
        self,
        paths: List[str],
        new_translations: List[float],
        new_rotation_eulers: List[float],
        new_rotation_orders: List[int],
        new_scales: List[float],
    ):
        """Function ran by _transform_selected_prims(). Can be overridden to change the behavior
        Do not remove this function: can be used outside of this code"""
        self._alert_if_selection_has_instance_proxies()
        omni.kit.commands.create(
            "TransformMultiPrimsSRTCpp",
            count=len(paths),
            no_undo=True,
            paths=paths,
            new_translations=new_translations,
            new_rotation_eulers=new_rotation_eulers,
            new_rotation_orders=new_rotation_orders,
            new_scales=new_scales,
            usd_context_name=self._usd_context_name,
            time_code=self._get_current_time_code().GetValue(),
        ).do()

    @carb.profiler.profile
    def _transform_all_selected_prims_to_manipulator_pivot(
        self,
        new_manipulator_transform: Gf.Matrix4d,
        dirty_ops: OpFlag,
    ):
        paths = []
        new_translations = []
        new_rotation_eulers = []
        new_rotation_orders = []
        new_scales = []
        old_translations = []
        old_rotation_eulers = []
        old_rotation_orders = []
        old_scales = []

        self._xform_cache.Clear()

        # any op may trigger a translation change if multi-manipulating
        should_update_translate = dirty_ops & OpFlag.TRANSLATE or len(self._xformable_prim_paths) > 1
        should_update_rotate = dirty_ops & OpFlag.ROTATE
        should_update_scale = dirty_ops & OpFlag.SCALE

        for path in self._xformable_prim_paths_sorted:
            if self._custom_manipulator_enabled and self._should_skip_custom_manipulator_path(path.pathString):
                continue

            selected_prim = self._stage.GetPrimAtPath(path)
            # We check whether path is in all_xformable_prim_data_curr because it may have not made it the dictionary if an error occured
            if not selected_prim or path not in self.all_xformable_prim_data_curr:
                continue

            (s, r, ro, t, selected_pivot) = self.all_xformable_prim_data_curr[path]

            selected_parent_to_world_mtx = self._xform_cache.GetParentToWorldTransform(selected_prim)
            world_to_parent_mtx = selected_parent_to_world_mtx.GetInverse()

            selected_local_mtx_new = new_manipulator_transform * world_to_parent_mtx * selected_pivot.GetInverse()

            if should_update_translate:
                translation = selected_local_mtx_new.ExtractTranslation()

            if should_update_rotate:
                # Construct the new rotation from old scale and translation.
                # Don't use Factor because it won't be able to tell if scale is positive or negative and can result in flipped rotation
                old_s_mtx = Gf.Matrix4d(1.0)
                old_s_mtx.SetScale(Gf.Vec3d(s))
                old_t_mtx = Gf.Matrix4d(1.0)
                old_t_mtx.SetTranslate(Gf.Vec3d(t))
                rot_new = (old_s_mtx.GetInverse() * selected_local_mtx_new * old_t_mtx.GetInverse()).ExtractRotation()

                axes = [Gf.Vec3d.XAxis(), Gf.Vec3d.YAxis(), Gf.Vec3d.ZAxis()]
                decomp_rot = rot_new.Decompose(axes[ro[2]], axes[ro[1]], axes[ro[0]])
                index_order = Gf.Vec3i()
                for i in range(3):
                    index_order[ro[i]] = 2 - i

                rotation = Gf.Vec3d(decomp_rot[index_order[0]], decomp_rot[index_order[1]], decomp_rot[index_order[2]])
                rotation = find_best_euler_angles(Gf.Vec3d(r), rotation, ro)

            if should_update_scale:
                # Construct the new scale from old rotation and translation.
                # Don't use Factor because it won't be able to tell if scale is positive or negative and can result in flipped rotation
                old_rt_mtx = compose_transform_ops_to_matrix(t, r, ro, Gf.Vec3d(1))
                new_s_mtx = selected_local_mtx_new * old_rt_mtx.GetInverse()
                scale = Gf.Vec3d(new_s_mtx[0][0], new_s_mtx[1][1], new_s_mtx[2][2])

            # any op may trigger a translation change if multi-manipulating
            translation = translation if should_update_translate else t
            rotation = rotation if should_update_rotate else r
            scale = scale if should_update_scale else s

            paths.append(path.pathString)
            new_translations += [translation[0], translation[1], translation[2]]
            new_rotation_eulers += [rotation[0], rotation[1], rotation[2]]
            new_rotation_orders += [ro[0], ro[1], ro[2]]
            new_scales += [scale[0], scale[1], scale[2]]
            old_translations += [t[0], t[1], t[2]]
            old_rotation_eulers += [r[0], r[1], r[2]]
            old_rotation_orders += [ro[0], ro[1], ro[2]]
            old_scales += [s[0], s[1], s[2]]

            xform_tuple = (scale, rotation, ro, translation, selected_pivot)
            self.all_xformable_prim_data_curr[path] = xform_tuple
            if path in self.consolidated_xformable_prim_data_curr:
                self.consolidated_xformable_prim_data_curr[path] = xform_tuple

        self._ignore_xform_data_change = True
        self._do_transform_all_selected_prims_to_manipulator_pivot(
            paths, new_translations, new_rotation_eulers, new_rotation_orders, new_scales
        )
        self._ignore_xform_data_change = False

    def _alert_if_selection_has_instance_proxies(self):
        if self._selected_instance_proxy_paths and (not self._warning_notification or self._warning_notification.dismissed):
            try:
                import omni.kit.notification_manager as nm
                self._warning_notification = nm.post_notification(
                    "Children of an instanced prim cannot be modified, uncheck Instanceable on the instanced prim to modify child prims.",
                    status=nm.NotificationStatus.WARNING
                )
            except ImportError:
                pass

    def _do_transform_all_selected_prims_to_manipulator_pivot(
        self,
        paths: List[str],
        new_translations: List[float],
        new_rotation_eulers: List[float],
        new_rotation_orders: List[int],
        new_scales: List[float],
    ):
        """Function ran by _transform_all_selected_prims_to_manipulator_pivot().
        Can be overridden to change the behavior. Do not remove this function: can be used outside of this code"""
        self._alert_if_selection_has_instance_proxies()
        omni.kit.commands.create(
            "TransformMultiPrimsSRTCpp",
            count=len(paths),
            no_undo=True,
            paths=paths,
            new_translations=new_translations,
            new_rotation_eulers=new_rotation_eulers,
            new_rotation_orders=new_rotation_orders,
            new_scales=new_scales,
            usd_context_name=self._usd_context_name,
            time_code=self._get_current_time_code().GetValue(),
        ).do()

    def _on_stage_event(self, event: carb.events.IEvent):
        if event.type == int(omni.usd.StageEventType.OPENED):
            self._on_stage_opened()
        elif event.type == int(omni.usd.StageEventType.CLOSING):
            self._on_stage_closing()

    def _on_stage_opened(self):
        self._stage = self._usd_context.get_stage()
        self._timeline_sub = self._timeline.get_timeline_event_stream().create_subscription_to_pop(
            self._on_timeline_event
        )
        self._current_time = self._timeline.get_current_time()
        self._xform_cache = UsdGeom.XformCache(self._get_current_time_code())

    def _on_stage_closing(self):
        self._stage = None
        self._xform_cache = None
        self._xformable_prim_paths.clear()
        self._xformable_prim_paths_sorted.clear()
        self._xformable_prim_paths_set.clear()
        self._xformable_prim_paths_prefix_set.clear()
        self._consolidated_xformable_prim_paths.clear()
        self._pivot_prim = None
        self._timeline_sub = None
        if self._stage_listener:
            self._stage_listener.Revoke()
            self._stage_listener = None
            carb.log_info("Tf.Notice.Revoke in PrimTransformModel")
        self._pending_changed_paths.clear()
        if self._update_prim_xform_from_prim_task_or_future is not None:
            self._update_prim_xform_from_prim_task_or_future.cancel()
            self._update_prim_xform_from_prim_task_or_future = None

    def on_selection_changed(self, selection: List[Sdf.Path]):
        if self._update_prim_xform_from_prim_task_or_future is not None:
            self._update_prim_xform_from_prim_task_or_future.cancel()
            self._update_prim_xform_from_prim_task_or_future = None

        self._selected_instance_proxy_paths.clear()
        self._xformable_prim_paths.clear()
        self._xformable_prim_paths_set.clear()
        self._xformable_prim_paths_prefix_set.clear()
        self._consolidated_xformable_prim_paths.clear()
        self._pivot_prim = None

        for sdf_path in selection:
            prim = self._stage.GetPrimAtPath(sdf_path)
            if prim and prim.IsA(UsdGeom.Xformable) and prim.IsActive():
                self._xformable_prim_paths.append(sdf_path)

        if self._xformable_prim_paths:
            # Make a sorted list so parents always appears before child
            self._xformable_prim_paths_sorted = self._xformable_prim_paths.copy()
            self._xformable_prim_paths_sorted.sort()

            # Find the most recently selected valid xformable prim as the pivot prim where the transform gizmo is located at.
            self._pivot_prim = self._stage.GetPrimAtPath(self._xformable_prim_paths[-1])

            # Get least common prims ancestors.
            # We do this so that if one selected prim is a descendant of other selected prim, the descendant prim won't be
            # transformed twice.
            self._consolidated_xformable_prim_paths = Sdf.Path.RemoveDescendentPaths(self._xformable_prim_paths)

            # Filter all instance proxy paths.
            for path in self._consolidated_xformable_prim_paths:
                prim = self._stage.GetPrimAtPath(path)
                if prim.IsInstanceProxy():
                    self._selected_instance_proxy_paths.add(sdf_path)

        self._xformable_prim_paths_set.update(self._xformable_prim_paths)
        for path in self._xformable_prim_paths_set:
            self._xformable_prim_paths_prefix_set.update(path.GetPrefixes())

        if self._update_transform_from_prims():
            self._item_changed(self._transform_item)

        # Happens when host widget is already enabled and first selection in a new stage
        if self._enabled_hosting_widget_count > 0 and self._stage_listener is None:
            self._stage_listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, self._on_objects_changed, self._stage)
            carb.log_info("Tf.Notice.Register in PrimTransformModel")

    def _should_keep_manipulator_orientation_unchanged(self, mode: str) -> bool:
        # Exclude snap_to_face. During snap_to_face operation, it may modify the orientation of object to confrom to surface
        # normal and the `new_manipulator_transform` param for `_transform_selected_prims` is set to the final transform
        # of the manipulated prim. However, if we use old rotation in the condition below, _no_scale_transform_manipulator
        # will not confrom to the new orientation, and _transform_selected_prims would double rotate the prims because it
        # sees the rotation diff between the old prim orientation (captured at on_began) vs new normal orient, instead of
        # current prim orientation vs new normal orientation.
        # Plus, it is nice to see the normal of the object changing while snapping.
        snap_provider_enabled = self.snap_settings_listener.snap_enabled and self.snap_settings_listener.snap_provider

        # When the manipulator is being manipulated as local translate or scale, we do not want to change the rotation of
        # the manipulator even if it's rotated, otherwise the direction of moving or scaling will change and can be very hard to control.
        # It can happen when you move a prim that has a constraint on it (e.g. lookAt)
        # In this case keep the rotation the same as on_began
        return (
            mode == c.TRANSFORM_MODE_LOCAL
            and (self._current_editing_op == Operation.TRANSLATE or self._current_editing_op == Operation.SCALE)
            and not snap_provider_enabled
        )

    def _should_keep_manipulator_translation_unchanged(self, mode: str) -> bool:
        # When the pivot placement is BBOX_CENTER and multiple prims being rotated, the bbox center may shifts, and the
        # rotation center will shift with them. This causes weird user experience. So we pin the rotation center until
        # mouse is released.
        return (
            self._current_editing_op == Operation.ROTATE
            and (self._placement == Placement.BBOX_CENTER or self._placement == Placement.BBOX_BASE)
        )

    def _get_transform_mode_for_current_op(self) -> str:
        mode = c.TRANSFORM_MODE_LOCAL
        if self._op_settings_listener.selected_op == c.TRANSFORM_OP_ROTATE:
            mode = self._op_settings_listener.rotation_mode
        elif self._op_settings_listener.selected_op == c.TRANSFORM_OP_MOVE:
            mode = self._op_settings_listener.translation_mode

        return mode

    # Adds a delay to the visual update during translate (only) manipulation
    # It's due to the renderer having a delay of rendering the mesh and the manipulator appears to drift apart.
    # It's only an estimation and may varying from scene/renderer setup.
    async def _delay_dirty(self, transform, id):
        if self._viewport_fps:
            render_frame_time = 1.0 / self._viewport_fps * self._settings.get(TRANSLATE_DELAY_FRAME_SETTING)
            while True:
                dt = await self._app.next_update_async()
                render_frame_time -= dt
                # break a frame early
                if render_frame_time < dt:
                    break

        # cancel earlier job if a later one catches up (fps suddenly changed?)
        earlier_tasks_or_futures = []
        for key, task_or_future in self._delay_dirty_tasks_or_futures.items():
            if key < id:
                earlier_tasks_or_futures.append(key)
                task_or_future.cancel()
            else:
                break
        for key in earlier_tasks_or_futures:
            self._delay_dirty_tasks_or_futures.pop(key)

        self._transform = transform
        self._item_changed(self._transform_item)
        self._delay_dirty_tasks_or_futures.pop(id)

    def _update_temp_pivot_world_position(self):
        if type(self._transform) is not list:
            return
        new_world_position = self._transform[12:15]
        self._settings.set_float_array(
            self._selection_pivot_position_path,
            new_world_position
        )

    @carb.profiler.profile
    def _update_transform_from_prims(self):
        xform_flattened = self._calculate_transform_from_prim()
        if self._transform != xform_flattened:
            self._transform = xform_flattened

            # update setting with new pivot placement position
            self._update_temp_pivot_world_position()

            return True
        return False

    def _calculate_transform_from_prim(self):
        if not self._stage:
            return False

        if not self._pivot_prim:
            return False

        self._xform_cache.Clear()
        cur_time = self._get_current_time_code()
        mode = self._get_transform_mode_for_current_op()

        pivot_inv = get_local_transform_pivot_inv(self._pivot_prim, cur_time)
        if self._should_keep_manipulator_orientation_unchanged(mode):
            pivot_prim_path = self._pivot_prim.GetPath()

            (s, r, ro, t) = omni.usd.get_local_transform_SRT(self._pivot_prim, cur_time)
            pivot = get_local_transform_pivot_inv(self._pivot_prim, self._get_current_time_code()).GetInverse()
            self.all_xformable_prim_data_curr[pivot_prim_path] = (s, r, ro, t) + (pivot,)

            # This method may be called from _on_op_listener_changed, before any gesture has started
            # in which case _all_xformable_prim_data_prev would be empty
            piv_xf_tuple = self._all_xformable_prim_data_prev.get(pivot_prim_path, False)
            if not piv_xf_tuple:
                piv_xf_tuple = omni.usd.get_local_transform_SRT(self._pivot_prim, cur_time)
                pv_xf_pivot = get_local_transform_pivot_inv(
                    self._pivot_prim, self._get_current_time_code()
                ).GetInverse()
                self._all_xformable_prim_data_prev[self._pivot_prim.GetPath()] = piv_xf_tuple + (pv_xf_pivot,)
            (s_p, r_p, ro_p, t_p, t_piv) = piv_xf_tuple

            xform = self._construct_transform_matrix_from_SRT(t, r_p, ro_p, s, pivot_inv)
            parent = self._xform_cache.GetLocalToWorldTransform(self._pivot_prim.GetParent())
            xform *= parent
        else:
            xform = self._xform_cache.GetLocalToWorldTransform(self._pivot_prim)
        xform = pivot_inv.GetInverse() * xform

        if self._should_keep_manipulator_translation_unchanged(mode):
            xform.SetTranslateOnly((self._transform[12], self._transform[13], self._transform[14]))
        else:
            # if there's only one selection, we always use LAST_PRIM_PIVOT though
            if (
                self._placement != Placement.LAST_PRIM_PIVOT
                and self._placement != Placement.REF_PRIM
            ):
                average_translation = Gf.Vec3d(0.0)
                if self._placement == Placement.BBOX_CENTER or self._placement == Placement.BBOX_BASE:
                    world_bound = Gf.Range3d()

                def get_prim_translation(xformable):
                    xformable_world_mtx = self._xform_cache.GetLocalToWorldTransform(xformable)
                    xformable_pivot_inv = get_local_transform_pivot_inv(xformable, cur_time)
                    xformable_world_mtx = xformable_pivot_inv.GetInverse() * xformable_world_mtx
                    return xformable_world_mtx.ExtractTranslation()

                for path in self._xformable_prim_paths:
                    xformable = self._stage.GetPrimAtPath(path)
                    if self._placement == Placement.SELECTION_CENTER:
                        average_translation += get_prim_translation(xformable)
                    elif self._placement == Placement.BBOX_CENTER or Placement.BBOX_BASE:
                        bound_range = self._usd_context.compute_path_world_bounding_box(path.pathString)
                        bound_range = Gf.Range3d(Gf.Vec3d(*bound_range[0]), Gf.Vec3d(*bound_range[1]))
                        if not bound_range.IsEmpty():
                            world_bound = Gf.Range3d.GetUnion(world_bound, bound_range)
                        else:
                            # extend world bound with tranlation for prims with zero bbox, e.g. Xform, Camera
                            prim_translation = get_prim_translation(xformable)
                            world_bound.UnionWith(prim_translation)

                if self._placement == Placement.SELECTION_CENTER:
                    average_translation /= len(self._xformable_prim_paths)
                elif self._placement == Placement.BBOX_CENTER:
                    average_translation = world_bound.GetMidpoint()
                elif self._placement == Placement.BBOX_BASE:
                    # xform may not have bbox but its descendants may have, exclude cases that only xform are selected
                    if not world_bound.IsEmpty():
                        bbox_center = world_bound.GetMidpoint()
                        bbox_size = world_bound.GetSize()

                        if UsdGeom.GetStageUpAxis(self._stage) == UsdGeom.Tokens.y:
                            # Y-up world
                            average_translation = bbox_center - Gf.Vec3d(0.0, bbox_size[1] / 2.0, 0.0)
                        else:
                            # Z-up world
                            average_translation = bbox_center - Gf.Vec3d(0.0, 0.0, bbox_size[2] / 2.0)
                    else:
                        # fallback to SELECTION_CENTER
                        average_translation /= len(self._xformable_prim_paths)

                # Only take the translate from selected prim average.
                # The rotation and scale still comes from pivot prim
                xform.SetTranslateOnly(average_translation)

        # instead of using RemoveScaleShear, additional steps made to handle negative scale properly
        scale, _, _, _ = omni.usd.get_local_transform_SRT(self._pivot_prim, cur_time)

        scale_epsilon = 1e-6
        for i in range(3):
            if Gf.IsClose(scale[i], 0.0, scale_epsilon):
                scale[i] = -scale_epsilon if scale[i] < 0 else scale_epsilon

        inverse_scale = Gf.Matrix4d().SetScale(Gf.Vec3d(1.0 / scale[0], 1.0 / scale[1], 1.0 / scale[2]))
        xform = inverse_scale * xform
        # this is the average xform without scale
        self._no_scale_transform_manipulator = Gf.Matrix4d(xform)
        # store the scale separately
        self._scale_manipulator = Gf.Vec3d(scale)

        # Visual transform of the manipulator
        xform = xform.RemoveScaleShear()

        if mode == c.TRANSFORM_MODE_GLOBAL:
            xform = xform.SetTranslate(xform.ExtractTranslation())

        return flatten(xform)

    def _construct_transform_matrix_from_SRT(
        self,
        translation: Gf.Vec3d,
        rotation_euler: Gf.Vec3d,
        rotation_order: Gf.Vec3i,
        scale: Gf.Vec3d,
        pivot_inv: Gf.Matrix4d,
    ):
        trans_mtx = Gf.Matrix4d()
        rot_mtx = Gf.Matrix4d()
        scale_mtx = Gf.Matrix4d()

        trans_mtx.SetTranslate(Gf.Vec3d(translation))

        axes = [Gf.Vec3d.XAxis(), Gf.Vec3d.YAxis(), Gf.Vec3d.ZAxis()]
        rotation = (
            Gf.Rotation(axes[rotation_order[0]], rotation_euler[rotation_order[0]])
            * Gf.Rotation(axes[rotation_order[1]], rotation_euler[rotation_order[1]])
            * Gf.Rotation(axes[rotation_order[2]], rotation_euler[rotation_order[2]])
        )
        rot_mtx.SetRotate(rotation)
        scale_mtx.SetScale(Gf.Vec3d(scale))
        return pivot_inv * scale_mtx * rot_mtx * pivot_inv.GetInverse() * trans_mtx

    def _on_op_listener_changed(self, type: OpSettingsListener.CallbackType, value: str):
        if type == OpSettingsListener.CallbackType.OP_CHANGED:
            # cancel all delayed tasks
            for task_or_future in self._delay_dirty_tasks_or_futures.values():
                task_or_future.cancel()
            self._delay_dirty_tasks_or_futures.clear()

            self._update_transform_from_prims()
            self._item_changed(self._transform_item)
        elif type == OpSettingsListener.CallbackType.TRANSLATION_MODE_CHANGED:
            if self._op_settings_listener.selected_op == c.TRANSFORM_OP_MOVE:
                if self._update_transform_from_prims():
                    self._item_changed(self._transform_item)
        elif type == OpSettingsListener.CallbackType.ROTATION_MODE_CHANGED:
            if self._op_settings_listener.selected_op == c.TRANSFORM_OP_ROTATE:
                if self._update_transform_from_prims():
                    self._item_changed(self._transform_item)

    def _update_placement(self, placement_str: str):
        if placement_str == prim_c.MANIPULATOR_PLACEMENT_SELECTION_CENTER:
            placement = Placement.SELECTION_CENTER
        elif placement_str == prim_c.MANIPULATOR_PLACEMENT_BBOX_CENTER:
            placement = Placement.BBOX_CENTER
        elif placement_str == prim_c.MANIPULATOR_PLACEMENT_PICK_REF_PRIM:
            placement = Placement.REF_PRIM
        elif placement_str == prim_c.MANIPULATOR_PLACEMENT_BBOX_BASE:
            placement = Placement.BBOX_BASE
        else:  # placement == prim_c.MANIPULATOR_PLACEMENT_LAST_PRIM_PIVOT or bad values
            placement = Placement.LAST_PRIM_PIVOT

        if placement != self._placement:
            if placement == Placement.LAST_PRIM_PIVOT and self._placement == Placement.REF_PRIM:
                # reset the pivot prim in case it was changed by MANIPULATOR_PLACEMENT_PICK_REF_PRIM
                if self._xformable_prim_paths:
                    self._pivot_prim = self._stage.GetPrimAtPath(self._xformable_prim_paths[-1])

            self._placement = placement
            if self._update_transform_from_prims():
                self._item_changed(self._transform_item)

    def _on_placement_setting_changed(self, item, event_type):
        placement_str = self._dict.get(item)
        self._update_placement(placement_str)

    def _check_update_selected_instance_proxy_list(self, path: Sdf.Path, resynced):
        def track_or_remove_from_instance_proxy_list(prim):
            valid_proxy = prim and prim.IsActive() and prim.IsInstanceProxy()
            if valid_proxy:
                self._selected_instance_proxy_paths.add(prim.GetPath())
            else:
                self._selected_instance_proxy_paths.discard(prim.GetPath())

        prim_path = path.GetPrimPath()
        changed_prim = self._stage.GetPrimAtPath(prim_path)

        # Update list of instance proxy paths.
        if resynced and path.IsPrimPath():
            if prim_path in self._consolidated_xformable_prim_paths:
                # Quick path if it's selected already.
                track_or_remove_from_instance_proxy_list(changed_prim)
            else:
                # Slow path to verify if any of its ancestors are changed.
                for path in self._consolidated_xformable_prim_paths:
                    if not path.HasPrefix(prim_path):
                        continue

                    prim = self._stage.GetPrimAtPath(path)
                    track_or_remove_from_instance_proxy_list(prim)

    @carb.profiler.profile
    async def _update_transform_from_prims_async(self):
        try:
            check_all_prims = (
                self._placement != Placement.LAST_PRIM_PIVOT
                and self._placement != Placement.REF_PRIM
                and len(self._xformable_prim_paths) > 1
            )
            pivot_prim_path = self._pivot_prim.GetPath()

            for p, resynced in self._pending_changed_paths.items():
                self._check_update_selected_instance_proxy_list(p, resynced)

                prim_path = p.GetPrimPath()
                # Update either check_all_prims
                #        or prim_path is a prefix of pivot_prim_path (pivot prim's parent affect pivot prim transform)
                # Note: If you move the manipulator and the prim flies away while manipulator stays in place, check this
                #       condition!
                if (
                    # check _xformable_prim_paths_prefix_set so that if the parent path of selected prim(s) changed, it
                    # can still update manipulator transform
                    prim_path in self._xformable_prim_paths_prefix_set
                    if check_all_prims
                    else pivot_prim_path.HasPrefix(prim_path)
                ):
                    if self._path_may_affect_transform(p):
                        # only delay the visual update in translate mode.
                        should_delay_frame = self._settings.get(TRANSLATE_DELAY_FRAME_SETTING) > 0
                        if self._op_settings_listener.selected_op == c.TRANSFORM_OP_MOVE and should_delay_frame:
                            xform = self._calculate_transform_from_prim()
                            id = self._app.get_update_number()
                            self._delay_dirty_tasks_or_futures[id] = run_coroutine(self._delay_dirty(xform, id))
                        else:
                            if self._update_transform_from_prims():
                                self._item_changed(self._transform_item)
                        break
        except Exception as e:
            carb.log_error(traceback.format_exc())
        finally:
            self._pending_changed_paths.clear()
            self._update_prim_xform_from_prim_task_or_future = None

    @carb.profiler.profile
    def _on_objects_changed(self, notice, sender):
        if not self._pivot_prim:
            return

        # collect resynced paths so that removed/added xformOps triggers refresh
        for path in notice.GetResyncedPaths():
            self._pending_changed_paths[path] = True

        # collect changed only paths
        for path in notice.GetChangedInfoOnlyPaths():
            self._pending_changed_paths[path] = False

        # if an operation is in progess, record all dirty xform path
        if self._current_editing_op is not None and not self._ignore_xform_data_change:
            self._pending_changed_paths_for_xform_data.update(notice.GetChangedInfoOnlyPaths())

        if self._update_prim_xform_from_prim_task_or_future is None or self._update_prim_xform_from_prim_task_or_future.done():
            self._update_prim_xform_from_prim_task_or_future = run_coroutine(self._update_transform_from_prims_async())

    def _on_timeline_event(self, e: carb.events.IEvent):
        if e.type == int(omni.timeline.TimelineEventType.CURRENT_TIME_TICKED):
            current_time = e.payload["currentTime"]
            if current_time != self._current_time:
                self._current_time = current_time
                self._xform_cache.SetTime(self._get_current_time_code())
                # TODO only update transform if this prim or ancestors transforms are time varying?
                if self._update_transform_from_prims():
                    self._item_changed(self._transform_item)

    def _get_current_time_code(self):
        return Usd.TimeCode(omni.usd.get_frame_time_code(self._current_time, self._stage.GetTimeCodesPerSecond()))

    def _set_default_settings(self):
        self._settings.set_default(TRANSFORM_GIZMO_CUSTOM_MANIPULATOR_ENABLED, False)
        self._settings.set_default(TRANSFORM_GIZMO_IS_USING, False)
        self._settings.set_default(TRANSFORM_GIZMO_TRANSLATE_DELTA_XYZ, [0, 0, 0])
        self._settings.set_default(TRANSFORM_GIZMO_ROTATE_DELTA_XYZW, [0, 0, 0, 1])
        self._settings.set_default(TRANSFORM_GIZMO_SCALE_DELTA_XYZ, [0, 0, 0])

    def _should_skip_custom_manipulator_path(self, path: str) -> bool:
        custom_manipulator_path_prims_settings_path = TRANSFORM_GIZMO_CUSTOM_MANIPULATOR_PRIMS + path
        return self._settings.get(custom_manipulator_path_prims_settings_path)

    def _path_may_affect_transform(self, path: Sdf.Path) -> bool:
        # Batched changes sent in a SdfChangeBlock may not have property name but only the prim path
        return not path.ContainsPropertyElements() or UsdGeom.Xformable.IsTransformationAffectedByAttrNamed(path.name)
