# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import math
from typing import List

import carb
import carb.profiler
import carb.settings
from pxr import Gf, Usd, UsdGeom


@carb.profiler.profile
def flatten(transform):
    """Convert array[4][4] to array[16]"""

    # flatten the matrix by hand
    # USING LIST COMPREHENSION IS VERY SLOW (e.g. return [item for sublist in transform for item in sublist]), which takes around 10ms.
    m0, m1, m2, m3 = transform[0], transform[1], transform[2], transform[3]
    return [
        m0[0],
        m0[1],
        m0[2],
        m0[3],
        m1[0],
        m1[1],
        m1[2],
        m1[3],
        m2[0],
        m2[1],
        m2[2],
        m2[3],
        m3[0],
        m3[1],
        m3[2],
        m3[3],
    ]


@carb.profiler.profile
def get_local_transform_pivot_inv(prim: Usd.Prim, time: Usd.TimeCode = Usd.TimeCode):
    xform = UsdGeom.Xformable(prim)
    xform_ops = xform.GetOrderedXformOps()
    if len(xform_ops):
        pivot_op_inv = xform_ops[-1]
        if (
            pivot_op_inv.GetOpType() == UsdGeom.XformOp.TypeTranslate
            and pivot_op_inv.IsInverseOp()
            and pivot_op_inv.GetName().endswith("pivot")
        ):
            return pivot_op_inv.GetOpTransform(time)
    return Gf.Matrix4d(1.0)


def compose_transform_ops_to_matrix(
    translation: Gf.Vec3d, rotation: Gf.Vec3d, rotation_order: Gf.Vec3i, scale: Gf.Vec3d
):
    axes = [Gf.Vec3d.XAxis(), Gf.Vec3d.YAxis(), Gf.Vec3d.ZAxis()]

    rot = []
    for i in range(3):
        axis_idx = rotation_order[i]
        rot.append(Gf.Rotation(axes[axis_idx], rotation[axis_idx]))

    rotation_mtx = Gf.Matrix4d(1.0)
    rotation_mtx.SetRotate(rot[0] * rot[1] * rot[2])

    valid_scale = Gf.Vec3d(0.0)
    for i in range(3):
        if abs(scale[i]) == 0:
            valid_scale[i] = 0.001
        else:
            valid_scale[i] = scale[i]

    scale_mtx = Gf.Matrix4d(1.0)
    scale_mtx.SetScale(valid_scale)

    translate_mtx = Gf.Matrix4d(1.0)
    translate_mtx.SetTranslate(translation)

    return scale_mtx * rotation_mtx * translate_mtx


def repeat(t: float, length: float) -> float:
    return t - (math.floor(t / length) * length)


def generate_compatible_euler_angles(euler: Gf.Vec3d, rotation_order: Gf.Vec3i) -> List[Gf.Vec3d]:
    equal_eulers = [euler]

    mid_order = rotation_order[1]

    equal = Gf.Vec3d()
    for i in range(3):
        if i == mid_order:
            equal[i] = 180 - euler[i]
        else:
            equal[i] = euler[i] + 180

    equal_eulers.append(equal)

    for i in range(3):
        equal[i] -= 360

    equal_eulers.append(equal)

    return equal_eulers


def find_best_euler_angles(old_rot_vec: Gf.Vec3d, new_rot_vec: Gf.Vec3d, rotation_order: Gf.Vec3i) -> Gf.Vec3d:
    equal_eulers = generate_compatible_euler_angles(new_rot_vec, rotation_order)
    nearest_euler = None

    for euler in equal_eulers:
        for i in range(3):
            euler[i] = repeat(euler[i] - old_rot_vec[i] + 180.0, 360.0) + old_rot_vec[i] - 180.0

        if nearest_euler is None:
            nearest_euler = euler
        else:
            distance_1 = (nearest_euler - old_rot_vec).GetLength()
            distance_2 = (euler - old_rot_vec).GetLength()

            if distance_2 < distance_1:
                nearest_euler = euler

    return nearest_euler
