# Copyright (c) 2018-2020, 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.
#
__all__ = ["PeerUserCameraManipulatorModel"]

import carb
import omni.usd

from omni.ui import scene as sc
from pxr import Gf, UsdGeom, Usd, Sdf


class PeerUserCameraManipulatorModel(sc.AbstractManipulatorModel):
    """The model tracks the position of peer user's camera in a live session."""

    class PositionItem(sc.AbstractManipulatorItem):
        def __init__(self):
            super().__init__()
            self.value = [0, 0, 0]
    
    class NdcPositionItem(sc.AbstractManipulatorItem):
        def __init__(self):
            super().__init__()
            self.value = [0, 0, 0]
    
    class ViewportCameraPositionItem(sc.AbstractManipulatorItem):
        def __init__(self):
            super().__init__()
            self.value = [0, 0, 0]

    def __init__(self, scene_view, camera_path, usd_context):
        super().__init__()

        # Cached position. It's updated only when on_transform_changed is explicitly requested.
        self.position = PeerUserCameraManipulatorModel.PositionItem()
        
        self.ndc_position = PeerUserCameraManipulatorModel.NdcPositionItem()

        self.viewport_camera_position = PeerUserCameraManipulatorModel.ViewportCameraPositionItem()
        
        self.__scene_view = scene_view

        self.__camera_path = Sdf.Path(camera_path)

        self.__usd_context = usd_context

        self.__update_position = True

        self.__update_ndc_position = True
    
    def destroy(self):
        self.position = None
        self.ndc_position = None
        self.viewport_camera_position = None
        self.__usd_context = None
        self.__scene_view = None
    
    def set_tracked_camera(self, camera_path):
        camera_path = Sdf.Path(camera_path)
        if camera_path != self.__camera_path:
            self.__camera_path = camera_path
            self.on_world_translation_changed()

    def get_item(self, identifier):
        if identifier == "position":
            return self.position
        elif identifier == "ndc_position":
            return self.ndc_position
        elif identifier == "viewport_camera_position":
            return self.viewport_camera_position
        
        return None

    def get_as_floats(self, item):
        if item == self.position:
            if self.__update_position:
                self.position.value = self.get_camera_position()
                self.__update_position = False

            return self.position.value
        elif item == self.ndc_position:
            if self.__update_ndc_position:
                self.ndc_position.value = self._get_ndc_position()
                self.__update_ndc_position = False
            
            return self.ndc_position.value
        elif item == self.viewport_camera_position:
            return self.viewport_camera_position.value

        return [0.0, 0.0, 0.0]

    def set_floats(self, item, value):
        if item != self.position:
            return
        
        if value and len(value) > 3:
            value = [value[0], value[1], value[2]]

        self.__update_position = False
        self.__update_ndc_position = False
        if value and item.value != value:
            # Set directly to the item
            item.value = value

            self.ndc_position.value = self._get_ndc_position()

            # This makes the manipulator updated
            self._item_changed(item)
    
    def on_view_translation_changed(self, viewport_camera_position=None):
        if viewport_camera_position and len(viewport_camera_position) > 3:
            viewport_camera_position = [
                viewport_camera_position[0], viewport_camera_position[1], viewport_camera_position[2]
            ]

        self.__update_ndc_position = True
        if viewport_camera_position:
            self.viewport_camera_position.value = viewport_camera_position
        self._item_changed(self.ndc_position)

    def on_world_translation_changed(self, position=None):
        """When position is provided, it will not re-calculate position but use provided value instead."""
        # Position is changed
        if position:
            self.set_floats(self.position, position)
        else:
            self.__update_position = True
            self.__update_ndc_position = True
            self._item_changed(self.position)
    
    def _get_ndc_position(self):
        if self.__update_position:
            self.position.value = self.get_camera_position()
            self.__update_position = False

        ndc_pos = self.__scene_view.scene.transform_space(sc.Space.WORLD, sc.Space.NDC, [*self.position.value, 1.0])

        return ndc_pos

    def get_camera_position(self):
        """Returns realtime position of currently selected object"""
        stage = self.__usd_context.get_stage()
        prim = stage.GetPrimAtPath(self.__camera_path)
        if not prim:
            return [0.0, 0.0, 0.0]

        # May be empty (UsdGeom.Xform for example), so build a scene-scaled constant box
        world_units = UsdGeom.GetStageMetersPerUnit(stage)
        if Gf.IsClose(world_units, 0.0, 1e-6):
            world_units = 0.01

        world_matrix = omni.usd.get_world_transform_matrix(prim)
        ref_position = world_matrix.ExtractTranslation()
        extent = Gf.Vec3d(0.1 / world_units)    # 10cm by default
        bboxMin = ref_position - extent
        bboxMax = ref_position + extent
        position = [(bboxMin[0] + bboxMax[0]) * 0.5, bboxMax[1], (bboxMin[2] + bboxMax[2]) * 0.5]

        return position
