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

import math
import omni.ui as ui
import omni.kit.usd.layers as layers

from omni.ui import scene as sc
from omni.ui import color as cl
from omni.kit.viewport.utility import get_active_viewport_window


class PeerUserCameraManipulator(sc.Manipulator):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Must
        self.__usd_context = kwargs.get("usd_context")

        self.__user_info = kwargs.get("user_info")

        self.__label_position_transform = None

        self.__offscreen_position_transform = None

        self.__pointer_position_transform = None

        self.__pointer_arrow_rotation_transform = None

        self.__offscreen_pointer = None

        self.__offscreen_pointer_size = 64

        self.__arrow_x_translation = 46

    def destroy(self):
        self.__clear()

    def __clear(self):
        if self.__label_position_transform:
            self.__label_position_transform.clear()
            self.__label_position_transform = None

        if self.__offscreen_position_transform:
            self.__offscreen_position_transform.clear()
            self.__offscreen_position_transform = None

        if self.__pointer_arrow_rotation_transform:
            self.__pointer_arrow_rotation_transform.clear()
            self.__pointer_arrow_rotation_transform = None

        if self.__pointer_position_transform:
            self.__pointer_position_transform.clear()
            self.__pointer_position_transform = None

    def __is_out_of_viewport(self, ndc_pos):
        # Make sure its not clipped
        if (ndc_pos[2] > 1.0) or (ndc_pos[0] < -1) or (ndc_pos[0] > 1) or (ndc_pos[1] < -1) or (ndc_pos[1] > 1):
            return True

        return False

    def __get_pointer_screen_pos_and_angle(self, viewport_window, ndc_pos):
        # When it's in the back of camera view,
        # Draw the widget at the bottom of viewport.
        if ndc_pos[2] > 1.0:
            ndc_pos[1] = -1.0
            ndc_pos[0] = -ndc_pos[0]

        ndc_pos[0] = min(max(ndc_pos[0], -1.0), 1.0)
        ndc_pos[1] = min(max(ndc_pos[1], -1.0), 1.0)

        radians = math.atan2(ndc_pos[1], ndc_pos[0])
        pointer_angles = radians

        frame = viewport_window.frame
        offset_x = ndc_pos[0] * frame.computed_width
        offset_y = ndc_pos[1] * frame.computed_height
        if ndc_pos[0] == -1.0:
            offset_x += self.__offscreen_pointer_size
        elif ndc_pos[0] == 1.0:
            offset_x -= self.__offscreen_pointer_size

        if ndc_pos[1] == -1.0:
            offset_y += self.__offscreen_pointer_size
        elif ndc_pos[1] == 1.0:
            offset_y -= self.__offscreen_pointer_size

        return (offset_x, offset_y, 0.0), pointer_angles

    def on_viewport_camera_changed(self, position):
        if position:
            self.__viewport_camera_position = position

    def on_build(self):
        """Called when the model is changed and rebuilds the whole slider"""
        if not self.model:
            return

        user_info = self.__user_info
        color = ui.color(*user_info.user_color)

        # Initialization
        if self.__offscreen_pointer is None:
            ndc_position = self.model.get_as_floats(self.model.get_item("ndc_position"))
            self.__offscreen_pointer = self.__is_out_of_viewport(ndc_position)

        ndc_position = self.model.get_as_floats(self.model.get_item("ndc_position"))
        if not ndc_position:
            return

        viewport_window = get_active_viewport_window()
        if not viewport_window:
            return

        bound_camera_world_translation = self.model.get_as_floats(self.model.get_item("viewport_camera_position"))
        if bound_camera_world_translation is None:
            return

        screen_position, pointer_angle = self.__get_pointer_screen_pos_and_angle(viewport_window, ndc_position)

        self.__offscreen_position_transform = sc.Transform(
            transform=sc.Matrix44.get_translation_matrix(*bound_camera_world_translation)
        )

        with self.__offscreen_position_transform:
            with sc.Transform(look_at=sc.Transform.LookAt.CAMERA):
                # Make sure it's in front of camera.
                with sc.Transform(transform=sc.Matrix44.get_translation_matrix(0, 0, -10000.0)):
                    with sc.Transform(scale_to=sc.Space.SCREEN):
                        self.__pointer_position_transform = sc.Transform(
                            transform=sc.Matrix44.get_translation_matrix(*screen_position)
                        )
                        with self.__pointer_position_transform:
                            sc.Arc(radius=28, color=color)
                            sc.Label(
                                layers.get_short_user_name(user_info.user_name), color=cl.white,
                                alignment=ui.Alignment.CENTER, size=16
                            )

                            arrow_transform = sc.Matrix44.get_rotation_matrix(0, 0, pointer_angle)
                            arrow_transform *= sc.Matrix44.get_translation_matrix(*[self.__arrow_x_translation, 0, 0])
                            self.__pointer_arrow_rotation_transform = sc.Transform(transform=arrow_transform)
                            with self.__pointer_arrow_rotation_transform:
                                sc.PolygonMesh(
                                    positions=[
                                        [14, 0, 0],
                                        [-14, 14, 0],
                                        [-14, -14, 0]
                                    ],
                                    colors=[color] * 3,
                                    vertex_counts=[3],
                                    vertex_indices=[0, 1, 2]
                                )

        position = self.model.get_as_floats(self.model.get_item("position"))
        name_size = len(self.__user_info.user_name) * 16
        self.__label_position_transform = sc.Transform(transform=sc.Matrix44.get_translation_matrix(*position))
        with self.__label_position_transform:
            with sc.Transform(look_at=sc.Transform.LookAt.CAMERA):
                with sc.Transform(scale_to=sc.Space.SCREEN):
                    with sc.Transform(transform=sc.Matrix44.get_translation_matrix(0, 50, 0)):
                        sc.Rectangle(color=color, width=name_size, height=40)
                        sc.Label(user_info.user_name, color=cl.white, alignment=ui.Alignment.CENTER, size=16)

        if self.__offscreen_pointer:
            self.__offscreen_position_transform.visible = True
            self.__label_position_transform.visible = False
        else:
            self.__label_position_transform.visible = True
            self.__offscreen_position_transform.visible = False

    def __invalidate(self):
        # Regenerate the manipulator
        self.__clear()
        self.invalidate()

    def on_model_updated(self, item):
        if not self.__offscreen_position_transform or not item:
            self.__invalidate()
            return

        viewport_window = get_active_viewport_window()
        if not viewport_window:
            self.__invalidate()
            return

        position_item = self.model.get_item("position")
        ndc_position_item = self.model.get_item("ndc_position")
        if item != position_item and item != ndc_position_item:
            return

        ndc_pos = self.model.get_as_floats(ndc_position_item)
        out_of_viewport = self.__is_out_of_viewport(ndc_pos)
        mode_changed = False
        if out_of_viewport != self.__offscreen_pointer:
            mode_changed = True
            self.__offscreen_pointer = out_of_viewport
            if out_of_viewport:
                self.__offscreen_position_transform.visible = True
                self.__label_position_transform.visible = False
            else:
                self.__offscreen_position_transform.visible = False
                self.__label_position_transform.visible = True

        if out_of_viewport:
            bound_camera_world_translation = self.model.get_as_floats(self.model.get_item("viewport_camera_position"))
            if bound_camera_world_translation is None:
                self.__invalidate()
                return

            position = bound_camera_world_translation
            screen_position, pointer_angle = self.__get_pointer_screen_pos_and_angle(viewport_window, ndc_pos)
            pointer_transform = sc.Matrix44.get_translation_matrix(*screen_position)
            self.__pointer_position_transform.transform = pointer_transform

            arrow_transform = sc.Matrix44.get_rotation_matrix(0, 0, pointer_angle)
            arrow_transform *= sc.Matrix44.get_translation_matrix(*[self.__arrow_x_translation, 0, 0])
            self.__pointer_arrow_rotation_transform.transform = arrow_transform

            # FIXME: Don't use sc.Matrix44.get_translation_matrix(*position), as it has
            # weird perf issue, which I think it's because of data type conversion.
            self.__offscreen_position_transform.transform = sc.Matrix44.get_translation_matrix(
                *[position[0], position[1], position[2]]
            )
        elif mode_changed or item == position_item:
            position = self.model.get_as_floats(position_item)
            self.__label_position_transform.transform = sc.Matrix44.get_translation_matrix(
                *[position[0], position[1], position[2]]
            )
