# Copyright 2016 Pixar
#
# Licensed under the Apache License, Version 2.0 (the "Apache License")
# with the following modification; you may not use this file except in
# compliance with the Apache License and the following modification to it:
# Section 6. Trademarks. is deleted and replaced with:
#
# 6. Trademarks. This License does not grant permission to use the trade
#    names, trademarks, service marks, or product names of the Licensor
#    and its affiliates, except as required to comply with Section 4(c) of
#    the License and to reproduce the content of the NOTICE file.
#
# You may obtain a copy of the Apache License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the Apache License with the above modification is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the Apache License for the specific
# language governing permissions and limitations under the Apache License.
#

# Note: This is just pulled directly from USD and modified
# (we only support Gf.Quat classes, but not Gf.Quaternion, so
# these tests are updated to reflect that)
# https://github.com/PixarAnimationStudios/USD/blob/release/pxr/base/gf/testenv/testGfQuaternion.py

from __future__ import division

__copyright__ = "Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved."
__license__ = """
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
import sys
import unittest

##### Uncomment for better assert messages
# if 'unittest.util' in __import__('sys').modules:
#    # Show full diff in self.assertEqual.
#    __import__('sys').modules['unittest.util']._MAX_LENGTH = 999999999


def setUpModule():
    # pre-load USD before anything else to prevent crash
    # from weird carb conflict
    # FIXME ^^^
    try:
        from . import update_path_and_load_usd

        update_path_and_load_usd()
    except ImportError:
        # not needed in Kit
        pass


try:
    import omni.kit.test

    TestClass = omni.kit.test.AsyncTestCase
except ImportError:
    TestClass = unittest.TestCase


class TestGfQuaternion(TestClass):
    def test_Constructors(self):
        from usdrt import Gf

        # self.assertIsInstance(Gf.Quaternion(), Gf.Quaternion)
        # self.assertIsInstance(Gf.Quaternion(0), Gf.Quaternion)
        # self.assertIsInstance(Gf.Quaternion(1, Gf.Vec3d(1,1,1)), Gf.Quaternion)

        self.assertIsInstance(Gf.Quath(Gf.Quath()), Gf.Quath)
        self.assertIsInstance(Gf.Quatf(Gf.Quatf()), Gf.Quatf)
        self.assertIsInstance(Gf.Quatd(Gf.Quatd()), Gf.Quatd)
        # Testing conversions between Quat[h,f,d]
        self.assertIsInstance(Gf.Quath(Gf.Quatf()), Gf.Quath)
        self.assertIsInstance(Gf.Quath(Gf.Quatd()), Gf.Quath)
        self.assertIsInstance(Gf.Quatf(Gf.Quath()), Gf.Quatf)
        self.assertIsInstance(Gf.Quatf(Gf.Quatd()), Gf.Quatf)
        self.assertIsInstance(Gf.Quatd(Gf.Quath()), Gf.Quatd)
        self.assertIsInstance(Gf.Quatd(Gf.Quatf()), Gf.Quatd)

    def test_Properties(self):
        from usdrt import Gf

        # nv edit - use Quat classses instead of Quaternion
        for Quat, Vec in ((Gf.Quatd, Gf.Vec3d), (Gf.Quatf, Gf.Vec3f), (Gf.Quath, Gf.Vec3h)):
            q = Quat()
            q.real = 10
            self.assertEqual(q.real, 10)
            q.imaginary = Vec(1, 2, 3)
            self.assertEqual(q.imaginary, Vec(1, 2, 3))

    def test_Methods(self):
        from usdrt import Gf

        # nv edit - use Quat classses instead of Quaternion
        for Quat, Vec in ((Gf.Quatd, Gf.Vec3d), (Gf.Quatf, Gf.Vec3f), (Gf.Quath, Gf.Vec3h)):
            q = Quat()
            self.assertEqual(Quat.GetIdentity(), Quat(1, Vec()))

            self.assertTrue(
                Quat.GetIdentity().GetLength() == 1
                and Gf.IsClose(Quat(1, Vec(2, 3, 4)).GetLength(), 5.4772255750516612, 0.00001)
            )

            q = Quat(1, Vec(2, 3, 4)).GetNormalized()
            self.assertTrue(
                Gf.IsClose(q.real, 0.182574, 0.0001)
                and Gf.IsClose(q.imaginary, Vec(0.365148, 0.547723, 0.730297), 0.00001)
            )

            # nv edit - linalg does not support arbitrary epsilon here
            # q = Quat(1,Vec(2,3,4)).GetNormalized(10)
            # self.assertEqual(q, Quat.GetIdentity())

            q = Quat(1, Vec(2, 3, 4))
            q.Normalize()
            self.assertTrue(
                Gf.IsClose(q.real, 0.182574, 0.0001)
                and Gf.IsClose(q.imaginary, Vec(0.365148, 0.547723, 0.730297), 0.00001)
            )

            # nv edit - linalg does not support arbitrary epsilon here
            # q = Quat(1,Vec(2,3,4)).Normalize(10)
            # self.assertEqual(q, Quat.GetIdentity())

            if Quat == Gf.Quath:
                # FIXME - below test does not pass for Quath:
                # AssertionError: Gf.Quath(1.0, Gf.Vec3h(0.0, 0.0, 0.0)) != Gf.Quath(1.0, Gf.Vec3h(-0.0, -0.0, -0.0))
                continue

            q = Quat.GetIdentity()
            self.assertEqual(q, q.GetInverse())
            q = Quat(1, Vec(1, 2, 3))
            q.Normalize()
            (re, im) = (q.real, q.imaginary)
            self.assertTrue(
                Gf.IsClose(q.GetInverse().real, re, 0.00001) and Gf.IsClose(q.GetInverse().imaginary, -im, 0.00001)
            )

    def test_Operators(self):
        from usdrt import Gf

        # nv edit - use Quat classses instead of Quaternion
        for Quat, Vec in ((Gf.Quatd, Gf.Vec3d), (Gf.Quatf, Gf.Vec3f), (Gf.Quath, Gf.Vec3h)):
            q1 = Quat(1, Vec(2, 3, 4))
            q2 = Quat(1, Vec(2, 3, 4))
            self.assertEqual(q1, q2)
            self.assertFalse(q1 != q2)
            q2.real = 2
            self.assertTrue(q1 != q2)

            q = Quat(1, Vec(2, 3, 4)) * Quat.GetIdentity()
            self.assertEqual(q, Quat(1, Vec(2, 3, 4)))

            q = Quat(1, Vec(2, 3, 4))
            q *= Quat.GetIdentity()
            self.assertEqual(q, Quat(1, Vec(2, 3, 4)))

            q *= 10
            self.assertEqual(q, Quat(10, Vec(20, 30, 40)))
            q = q * 10
            self.assertEqual(q, Quat(100, Vec(200, 300, 400)))
            q = 10 * q
            self.assertEqual(q, Quat(1000, Vec(2000, 3000, 4000)))
            q /= 100
            self.assertEqual(q, Quat(10, Vec(20, 30, 40)))
            q = q / 10
            self.assertEqual(q, Quat(1, Vec(2, 3, 4)))

            q += q
            self.assertEqual(q, Quat(2, Vec(4, 6, 8)))

            q -= Quat(1, Vec(2, 3, 4))
            self.assertEqual(q, Quat(1, Vec(2, 3, 4)))

            q = q + q
            self.assertEqual(q, Quat(2, Vec(4, 6, 8)))

            q = q - Quat(1, Vec(2, 3, 4))
            self.assertEqual(q, Quat(1, Vec(2, 3, 4)))

            q = q * q
            self.assertEqual(q, Quat(-28, Vec(4, 6, 8)))

            q1 = Quat(1, Vec(2, 3, 4)).GetNormalized()
            q2 = Quat(4, Vec(3, 2, 1)).GetNormalized()
            self.assertEqual(Gf.Slerp(0, q1, q2), q1)
            # nv edit - these are close but not exact with our implementation
            eps = 0.00001 if Quat in [Gf.Quatd, Gf.Quatf] else 0.001
            # self.assertEqual(Gf.Slerp(1, q1, q2), q2)
            self.assertTrue(Gf.IsClose(Gf.Slerp(1, q1, q2), q2, eps))
            # self.assertEqual(Gf.Slerp(0.5, q1, q2), Quat(0.5, Vec(0.5, 0.5, 0.5)))
            self.assertTrue(Gf.IsClose(Gf.Slerp(0.5, q1, q2), Quat(0.5, Vec(0.5, 0.5, 0.5)), eps))

            # code coverage goodness
            q1 = Quat(0, Vec(1, 1, 1))
            q2 = Quat(0, Vec(-1, -1, -1))
            q = Gf.Slerp(0.5, q1, q2)
            self.assertTrue(Gf.IsClose(q.real, 0, 0.00001) and Gf.IsClose(q.imaginary, Vec(1, 1, 1), 0.00001))

            q1 = Quat(0, Vec(1, 1, 1))
            q2 = Quat(0, Vec(1, 1, 1))
            q = Gf.Slerp(0.5, q1, q2)
            self.assertTrue(Gf.IsClose(q.real, 0, 0.00001) and Gf.IsClose(q.imaginary, Vec(1, 1, 1), 0.00001))

            self.assertEqual(q, eval(repr(q)))

            self.assertTrue(len(str(Quat())))

        for quatType in (Gf.Quatd, Gf.Quatf, Gf.Quath):
            q1 = quatType(1, [2, 3, 4])
            q2 = quatType(2, [3, 4, 5])

            self.assertTrue(Gf.IsClose(Gf.Dot(q1, q2), 40, 0.00001))


if __name__ == "__main__":
    unittest.main()
