# 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 math
from typing import List

import carb
from omni.kit.manipulator.transform.gestures import (
    RotateChangedGesture,
    RotateDragGesturePayload,
    ScaleChangedGesture,
    ScaleDragGesturePayload,
    TransformDragGesturePayload,
    TranslateChangedGesture,
    TranslateDragGesturePayload,
)
from omni.ui import scene as sc

from .model import AbstractTransformManipulatorModel
from .types import Operation


class SimpleTranslateChangedGesture(TranslateChangedGesture):
    def on_changed(self):
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, TranslateDragGesturePayload):
            return

        model = self.sender.model
        item = self.gesture_payload.changing_item
        translated = self.gesture_payload.moved_delta
        translation = [a + b for a, b in zip(translated, model.get_as_floats(item))]

        model.set_floats(item, translation)


class SimpleRotateChangedGesture(RotateChangedGesture):
    def on_began(self):
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, RotateDragGesturePayload):
            return

        model = self.sender.model
        item = self.gesture_payload.changing_item
        self._begin_rotation = model.get_as_floats(item).copy()
        self._begin_rotation_matrix = sc.Matrix44.get_rotation_matrix(
            self._begin_rotation[0], self._begin_rotation[1], self._begin_rotation[2], True
        )

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

        model = self.sender.model
        item = self.gesture_payload.changing_item

        axis = self.gesture_payload.axis
        angle = self.gesture_payload.angle

        # always reset begin_rotation in local rotation mode
        if not model.global_mode:
            self._begin_rotation = model.get_as_floats(item).copy()
            self._begin_rotation_matrix = sc.Matrix44.get_rotation_matrix(
                self._begin_rotation[0], self._begin_rotation[1], self._begin_rotation[2], True
            )

        # convert to radian
        angle = angle / 180.0 * math.pi

        # calculate rotation matrix around axis for angle in radian
        matrix = [
            math.cos(angle) + axis[0] ** 2 * (1 - math.cos(angle)),
            axis[0] * axis[1] * (1 - math.cos(angle)) + axis[2] * math.sin(angle),
            axis[0] * axis[2] * (1 - math.cos(angle)) - axis[1] * math.sin(angle),
            0,
        ]
        matrix += [
            axis[0] * axis[1] * (1 - math.cos(angle)) - axis[2] * math.sin(angle),
            math.cos(angle) + axis[1] ** 2 * (1 - math.cos(angle)),
            axis[1] * axis[2] * (1 - math.cos(angle)) + axis[0] * math.sin(angle),
            0,
        ]
        matrix += [
            axis[0] * axis[2] * (1 - math.cos(angle)) + axis[1] * math.sin(angle),
            axis[1] * axis[2] * (1 - math.cos(angle)) - axis[0] * math.sin(angle),
            math.cos(angle) + axis[2] ** 2 * (1 - math.cos(angle)),
            0,
        ]

        matrix += [0, 0, 0, 1]
        matrix = sc.Matrix44(*matrix)  # each 4 elements in list is a COLUMN in a row-vector matrix!

        rotate = self._begin_rotation_matrix * matrix

        # decompose back to x y z euler angle
        sy = math.sqrt(rotate[0] ** 2 + rotate[1] ** 2)
        is_singular = sy < 10 ** -6
        if not is_singular:
            z = math.atan2(rotate[1], rotate[0])
            y = math.atan2(-rotate[2], sy)
            x = math.atan2(rotate[6], rotate[10])
        else:
            z = math.atan2(-rotate[9], rotate[5])
            y = math.atan2(-rotate[2], sy)
            x = 0

        x = x / math.pi * 180.0
        y = y / math.pi * 180.0
        z = z / math.pi * 180.0

        model.set_floats(item, [x, y, z])


class SimpleScaleChangedGesture(ScaleChangedGesture):
    def on_began(self):
        if not self.gesture_payload or not self.sender or not isinstance(self.gesture_payload, ScaleDragGesturePayload):
            return

        model = self.sender.model
        item = self.gesture_payload.changing_item
        self._begin_scale = model.get_as_floats(item).copy()

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

        model = self.sender.model
        item = self.gesture_payload.changing_item

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

        scale_delta = [scale * v for v in axis]
        scale_vec = [0, 0, 0]
        for i in range(3):
            scale_vec[i] = self._begin_scale[i] * scale_delta[i] if scale_delta[i] else self._scale[i]

        model.set_floats(item, [scale_vec[0], scale_vec[1], scale_vec[2]])


class SimpleTransformModel(AbstractTransformManipulatorModel):
    """
    SimpleTransformModel is a model that provides basic S/R/T transform data processing. You can subscribe to callback
    to get manipulated data. Rotation is stored in degree and applis in XYZ order.
    """

    def __init__(self):
        super().__init__()
        self._translation = [0, 0, 0]
        self._rotation = [0, 0, 0]
        self._scale = [1, 1, 1]
        self._transform = sc.Matrix44(1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1)
        self._op = Operation.TRANSLATE
        self._global_mode = True
        self._dirty = True

    @property
    def global_mode(self) -> bool:
        """
        The Global vs Local transform mode.
        """
        return self._global_mode

    @global_mode.setter
    def global_mode(self, value: bool):
        if self._global_mode != value:
            self._global_mode = value
            self._dirty = True  # global vs local mode returns different Transform matrix
            self._item_changed(self._transform_item)

    def set_floats(self, item: AbstractTransformManipulatorModel.OperationItem, value: List[float]):
        if item == self._translate_item:
            self._translation[0] = value[0]
            self._translation[1] = value[1]
            self._translation[2] = value[2]
            self._dirty = True
            self._item_changed(self._translate_item)
        elif item == self._rotate_item:
            self._rotation[0] = value[0]
            self._rotation[1] = value[1]
            self._rotation[2] = value[2]
            self._dirty = True
            self._item_changed(self._rotate_item)
        elif item == self._scale_item:
            self._scale[0] = value[0]
            self._scale[1] = value[1]
            self._scale[2] = value[2]
            self._dirty = True
            self._item_changed(self._scale_item)
        else:
            carb.log_warn(f"Unsupported item {item}")

    def get_as_floats(self, item: AbstractTransformManipulatorModel.OperationItem) -> List[float]:
        if item == self._translate_item:
            return self._translation
        elif item == self._rotate_item:
            return self._rotation
        elif item == self._scale_item:
            return self._scale
        elif item is None or item == self._transform_item:
            if self._dirty:
                # Scale is not put into the Transform matrix because we don't want the TransformManipulator itself to scale
                # For a "global" style rotation gizmo (where the gizmo doesn't rotate), we don't want to put the rotation into Transform either.
                if self._global_mode:
                    self._transform = sc.Matrix44.get_translation_matrix(
                        self._translation[0], self._translation[1], self._translation[2]
                    )
                else:
                    self._transform = sc.Matrix44.get_translation_matrix(
                        self._translation[0], self._translation[1], self._translation[2]
                    ) * sc.Matrix44.get_rotation_matrix(self._rotation[0], self._rotation[1], self._rotation[2], True)
                self._dirty = False

            return self._transform
        else:
            carb.log_warn(f"Unsupported item {item}")
            return None

    def get_operation(self) -> Operation:
        return self._op

    def set_operation(self, op: Operation):
        if self._op != op:
            self._op = op
            self._item_changed(self._transform_item)
