__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_mat_info():
    """Return list of tuples of class, scalar size, item count, and format string"""
    from usdrt import Gf

    mat_info = [
        (Gf.Matrix3d, 8, 9, "d"),
        (Gf.Matrix3f, 4, 9, "f"),
        (Gf.Matrix4d, 8, 16, "d"),
        (Gf.Matrix4f, 4, 16, "f"),
    ]

    return mat_info


def get_mat_numpy_types():
    import numpy

    import usdrt

    equivalent_types = {
        usdrt.Gf.Matrix3d: numpy.double,
        usdrt.Gf.Matrix3f: numpy.single,
        usdrt.Gf.Matrix4d: numpy.double,
        usdrt.Gf.Matrix4f: numpy.single,
    }

    return equivalent_types


class TestGfMatBuffer(TestRtScenegraph):
    """Test buffer protocol for usdrt.Gf.Matrix* classes"""

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

        test_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

        for Mat, size, count, fmt in mat_info:
            v = Mat(*test_values[:count])
            view = memoryview(v)
            self.assertEquals(view.itemsize, size)
            self.assertEquals(view.ndim, 2)
            self.assertEquals(view.shape, (math.sqrt(count), math.sqrt(count)))
            self.assertEquals(view.format, fmt)
            self.assertEquals(view.obj, Mat(*test_values[:count]))

    @tc_logger
    def test_buffer_init_from_pxr_type(self):
        """Validate initialization from equivalent Pixar types"""
        import pxr.Gf

        import usdrt

        mat_info = get_mat_info()

        equivalent_types = {
            usdrt.Gf.Matrix3d: pxr.Gf.Matrix3d,
            usdrt.Gf.Matrix3f: pxr.Gf.Matrix3f,
            usdrt.Gf.Matrix4d: pxr.Gf.Matrix4d,
            usdrt.Gf.Matrix4f: pxr.Gf.Matrix4f,
        }

        test_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

        for Mat, size, count, fmt in mat_info:

            test_pxr_type = equivalent_types[Mat](*test_values[:count])
            test_mat = Mat(test_pxr_type)

            self.assertEquals(test_mat, Mat(*test_values[:count]))

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

        mat_info = get_mat_info()

        equivalent_types = get_mat_numpy_types()

        test_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

        for Mat, size, count, fmt in mat_info:

            py_mat = []
            stride = int(math.sqrt(count))
            for i in range(stride):
                py_mat.append(test_values[i * stride : i * stride + stride])

            test_numpy_type = numpy.array(py_mat, dtype=equivalent_types[Mat])
            test_mat = Mat(test_numpy_type)

            self.assertEquals(test_mat, Mat(*test_values[:count]))

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

        mat_info = get_mat_info()

        equivalent_types = get_mat_numpy_types()

        test_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

        for Mat, size, count, fmt in mat_info:

            py_mat = []
            stride = int(math.sqrt(count))
            for i in range(stride):
                py_mat.append(test_values[i * stride : i * stride + stride])

            test_numpy_type = numpy.array(py_mat, dtype=numpy.int_)
            with self.assertRaises(ValueError):
                test_mat = Mat(test_numpy_type)

    @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 numpy

        mat_info = get_mat_info()

        equivalent_types = get_mat_numpy_types()

        test_values = [[0, 1], [2, 3]]

        for Mat, size, count, fmt in mat_info:
            test_numpy_type = numpy.array(test_values, dtype=equivalent_types[Mat])

            with self.assertRaises(ValueError):
                test_mat = Mat(test_numpy_type)

    @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

        mat_info = get_mat_info()

        equivalent_types = get_mat_numpy_types()

        test_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

        for Mat, size, count, fmt in mat_info:
            test_array = numpy.array(test_values[:count], dtype=equivalent_types[Mat])

            with self.assertRaises(ValueError):
                test_mat = Mat(test_array)


class TestGfMatrixGetSetItem(TestRtScenegraph):
    """Test single-item accessors for usdrt.Gf.Matrix* classes"""

    @tc_logger
    def test_get_item(self):
        """Validate GetArrayItem on GfMatrix"""
        from usdrt import Gf

        for Mat in [Gf.Matrix3d, Gf.Matrix3f, Gf.Matrix4d, Gf.Matrix4f]:
            items = Mat.dimension[0] * Mat.dimension[1]
            m = Mat(*[i for i in range(items)])

            for i in range(items):
                self.assertTrue(m.GetArrayItem(i) == i)

    @tc_logger
    def test_set_item(self):
        """Validate SetArrayItem on GfMatrix"""
        from usdrt import Gf

        for Mat in [Gf.Matrix3d, Gf.Matrix3f, Gf.Matrix4d, Gf.Matrix4f]:
            m = Mat()
            self.assertEqual(m, Mat(1))

            items = Mat.dimension[0] * Mat.dimension[1]
            expected = Mat(*[i for i in range(items)])

            for i in range(items):
                m.SetArrayItem(i, i)

            self.assertNotEqual(m, Mat(1))
            self.assertEqual(m, expected)


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

    @tc_logger
    def test_copy(self):
        """Test __copy__"""
        mat_info = get_mat_info()

        for Mat, size, count, fmt in mat_info:
            x = Mat()
            y = x
            self.assertTrue(y is x)

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

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

    @tc_logger
    def test_deepcopy(self):
        """Test __deepcopy__"""
        mat_info = get_mat_info()

        for Mat, size, count, fmt in mat_info:
            x = Mat()
            y = x
            self.assertTrue(y is x)

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

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


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

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

        test_vec = Gf.Vec4d(9, 8, 7, 6)
        test_matrix = Gf.Matrix4d(test_vec)  # Sets matrix to diagonal form (ith element on the diagonal set to vec[i])
        expected = "Gf.Matrix4d(9.0, 0.0, 0.0, 0.0,  0.0, 8.0, 0.0, 0.0,  0.0, 0.0, 7.0, 0.0,  0.0, 0.0, 0.0, 6.0)"

        self.assertEqual(repr(test_matrix), expected)


class TestGfMatrixSetLookAtExtras(TestRtScenegraph):
    """More tests to validate SetLookAt"""

    @tc_logger
    def test_set_look_at_extras(self):
        """Test more points to ensure behavior match to Pixar implementation"""
        from pxr import Gf

        from usdrt import Gf as RtGf

        target = [(0, 0, -20), (0, 0, -30), (0, 0, -40), (-10, 0, -40), (-20, 0, -40), (-30, 0, -40)]
        eye = [(0, 0, 0), (0, 1, 0), (0, 2, 0), (0, 3, 0), (1, 3, 0), (2, 3, 0), (3, 3, 0)]

        for eye_pos in eye:
            for target_pos in target:
                pxr_result = Gf.Matrix4d(1).SetLookAt(Gf.Vec3d(*eye_pos), Gf.Vec3d(*target_pos), Gf.Vec3d(0, 1, 0))
                rt_result = RtGf.Matrix4d(1).SetLookAt(
                    RtGf.Vec3d(*eye_pos), RtGf.Vec3d(*target_pos), RtGf.Vec3d(0, 1, 0)
                )

                self.assertTrue(RtGf.IsClose(RtGf.Matrix4d(pxr_result), rt_result, 0.0001))
