__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 omni.kit.test
import copy
import math
import os
import platform
import sys
import time
import unittest

from . import TestRtScenegraph, tc_logger, update_path_and_load_usd


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


def get_quat_info():
    """Return list of tuples of class, scalar size, and format string"""
    from usdrt import Gf

    quat_info = [
        (Gf.Quatd, 8, "d"),
        (Gf.Quatf, 4, "f"),
        (Gf.Quath, 2, "e"),
    ]

    return quat_info


def get_quat_numpy_types():
    import numpy

    import usdrt

    equivalent_types = {
        usdrt.Gf.Quatd: numpy.double,
        usdrt.Gf.Quatf: numpy.single,
        usdrt.Gf.Quath: numpy.half,
    }

    return equivalent_types


class TestGfQuatBuffer(TestRtScenegraph):
    """Test buffer protocol for usdrt.Gf.Quat* classes"""

    @tc_logger
    def test_buffer_inspect(self):
        """Use memoryview to validate that buffer protocol is implemented"""
        quat_info = get_quat_info()

        test_values = [9, 8, 7, 1]

        for Quat, size, fmt in quat_info:
            q = Quat(test_values[3], test_values[0], test_values[1], test_values[2])
            view = memoryview(q)
            self.assertEquals(view.itemsize, size)
            self.assertEquals(view.ndim, 1)
            self.assertEquals(view.shape, (4,))
            self.assertEquals(view.format, fmt)
            self.assertEquals(view.obj, Quat(test_values[3], test_values[0], test_values[1], test_values[2]))

    @tc_logger
    def test_buffer_init(self):
        """Validate initialization from an array using buffer protocol"""
        import array

        quat_info = get_quat_info()

        test_values = [9, 8, 7, 1]

        for Quat, size, fmt in quat_info:
            if fmt == "e":
                # GfHalf not supported with array.array
                continue

            test_array = array.array(fmt, test_values)
            test_quat = Quat(test_array)

            self.assertEquals(test_quat, Quat(test_values[3], test_values[0], test_values[1], test_values[2]))

    @tc_logger
    def test_buffer_init_from_numpy_type(self):
        """Test initialization from numpy types"""
        import numpy

        quat_info = get_quat_info()

        equivalent_types = get_quat_numpy_types()

        test_values = [9, 8, 7, 1]

        for Quat, size, fmt in quat_info:
            test_numpy_type = numpy.array(test_values, dtype=equivalent_types[Quat])
            test_quat = Quat(test_numpy_type)

            self.assertEquals(test_quat, Quat(test_values[3], test_values[0], test_values[1], test_values[2]))

    @tc_logger
    def test_buffer_init_wrong_type(self):
        """Verify that an exception is raised with initializing via buffer with different data type"""
        import array

        quat_info = get_quat_info()

        test_values = [9, 8, 7, 1]

        for Quat, size, fmt in quat_info:
            test_array = array.array("h", test_values)

            with self.assertRaises(ValueError):
                test_quat = Quat(test_array)

    @tc_logger
    def test_buffer_init_wrong_size(self):
        """Verify that an exception is raised with initializing via buffer with different data array length"""
        import array

        quat_info = get_quat_info()

        test_values = [9, 8, 7, 6, 1]

        for Quat, size, fmt in quat_info:
            if fmt == "e":
                # GfHalf not supported with array.array
                continue

            test_array = array.array(fmt, test_values)

            with self.assertRaises(ValueError):
                test_vec = Quat(test_array)

    @tc_logger
    def test_buffer_init_wrong_dimension(self):
        """Verify that an exception is raised with initializing via buffer with different data dimensions"""
        import numpy

        quat_info = get_quat_info()
        equivalent_types = get_quat_numpy_types()

        test_values = [9, 8, 7, 1]

        for Quat, size, fmt in quat_info:
            test_array = numpy.array([test_values, test_values], dtype=equivalent_types[Quat])

            with self.assertRaises(ValueError):
                test_quat = Quat(test_array)


class TestGfQuatGetSetItem(TestRtScenegraph):
    """Test item accessors for usdrt.Gf.Quat* classes"""

    @tc_logger
    def test_get_item(self):
        """Validate __getitem__ on GfQuat"""
        from usdrt import Gf

        for Quat in [Gf.Quatf, Gf.Quatd, Gf.Quath]:
            q = Quat(0.5, 1, 2, 3)

            # Note that these are in memory layout order,
            # which is different from constructor order
            self.assertEqual(q[0], 1)
            self.assertEqual(q[1], 2)
            self.assertEqual(q[2], 3)
            self.assertEqual(q[3], 0.5)
            self.assertEqual(q[-4], 1)
            self.assertEqual(q[-3], 2)
            self.assertEqual(q[-2], 3)
            self.assertEqual(q[-1], 0.5)

    @tc_logger
    def test_set_item(self):
        """Validate __setitem__ on GfQuat"""
        from usdrt import Gf

        for Quat in [Gf.Quatf, Gf.Quatd, Gf.Quath]:
            q = Quat()

            self.assertEqual(q.imaginary[0], 0)
            self.assertEqual(q.imaginary[1], 0)
            self.assertEqual(q.imaginary[2], 0)
            self.assertEqual(q.real, 0)

            q[0] = 1
            q[1] = 2
            q[2] = 3
            q[3] = 0.5

            self.assertEqual(q.imaginary[0], 1)
            self.assertEqual(q.imaginary[1], 2)
            self.assertEqual(q.imaginary[2], 3)
            self.assertEqual(q.real, 0.5)

    @tc_logger
    def test_get_item_slice(self):
        """Validate __getitem__ with a slice on GfQuat"""
        from usdrt import Gf

        for Quat in [Gf.Quatf, Gf.Quatd, Gf.Quath]:
            q = Quat(0.5, 1, 2, 3)

            self.assertEqual(q[1:], [2, 3, 0.5])

    @tc_logger
    def test_set_item_slice(self):
        """Validate __setitem__ with a slice on GfQuat"""
        from usdrt import Gf

        for Quat in [Gf.Quatf, Gf.Quatd, Gf.Quath]:
            q = Quat(0.5, 1, 2, 3)

            self.assertEqual(q.imaginary[0], 1)
            self.assertEqual(q.imaginary[1], 2)
            self.assertEqual(q.imaginary[2], 3)
            self.assertEqual(q.real, 0.5)

            q[1:] = [5, 6, 7]

            self.assertEqual(q.imaginary[0], 1)
            self.assertEqual(q.imaginary[1], 5)
            self.assertEqual(q.imaginary[2], 6)
            self.assertEqual(q.real, 7)


class TestGfQuatCopy(TestRtScenegraph):
    """Test copy and deepcopy support for Gf.Quat* classes"""

    @tc_logger
    def test_copy(self):
        """Test __copy__"""
        quat_info = get_quat_info()

        for Quat, size, fmt in quat_info:
            x = Quat()
            y = x
            self.assertTrue(y is x)

            y[0] = 10
            self.assertEqual(x[0], 10)

            z = copy.copy(y)
            self.assertFalse(z is y)
            y[0] = 100
            self.assertEqual(z[0], 10)

    @tc_logger
    def test_deepcopy(self):
        """Test __deepcopy__"""
        quat_info = get_quat_info()

        for Quat, size, fmt in quat_info:
            x = Quat()
            y = x
            self.assertTrue(y is x)

            y[0] = 10
            self.assertEqual(x[0], 10)

            z = copy.deepcopy(y)
            self.assertFalse(z is y)
            y[0] = 100
            self.assertEqual(z[0], 10)


class TestGfQuatDebugging(TestRtScenegraph):
    """Test human readable output functions for debugging"""

    @tc_logger
    def test_repr_representation(self):
        """Test __repr__"""
        from usdrt import Gf

        test_quat = Gf.Quatf(0.5, 1, 2, 3)
        expected = "Gf.Quatf(0.5, Gf.Vec3f(1.0, 2.0, 3.0))"

        self.assertEqual(repr(test_quat), expected)
