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

    vec_info = [
        (Gf.Vec2d, 8, 2, "d"),
        (Gf.Vec2f, 4, 2, "f"),
        (Gf.Vec2h, 2, 2, "e"),
        (Gf.Vec2i, 4, 2, "i"),
        (Gf.Vec3d, 8, 3, "d"),
        (Gf.Vec3f, 4, 3, "f"),
        (Gf.Vec3h, 2, 3, "e"),
        (Gf.Vec3i, 4, 3, "i"),
        (Gf.Vec4d, 8, 4, "d"),
        (Gf.Vec4f, 4, 4, "f"),
        (Gf.Vec4h, 2, 4, "e"),
        (Gf.Vec4i, 4, 4, "i"),
    ]

    return vec_info


def get_vec_numpy_types():
    import numpy

    import usdrt

    equivalent_types = {
        usdrt.Gf.Vec2d: numpy.double,
        usdrt.Gf.Vec2f: numpy.single,
        usdrt.Gf.Vec2h: numpy.half,
        usdrt.Gf.Vec2i: numpy.int,
        usdrt.Gf.Vec3d: numpy.double,
        usdrt.Gf.Vec3f: numpy.single,
        usdrt.Gf.Vec3h: numpy.half,
        usdrt.Gf.Vec3i: numpy.int,
        usdrt.Gf.Vec4d: numpy.double,
        usdrt.Gf.Vec4f: numpy.single,
        usdrt.Gf.Vec4h: numpy.half,
        usdrt.Gf.Vec4i: numpy.int,
    }

    return equivalent_types


class TestGfVecBuffer(TestRtScenegraph):
    """Test buffer protocol for usdrt.Gf.Vec* classes"""

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

        vec_info = get_vec_info()

        test_values = [9, 8, 7, 6]

        for Vec, size, count, fmt in vec_info:
            v = Vec(*test_values[:count])
            view = memoryview(v)
            self.assertEquals(view.itemsize, size)
            self.assertEquals(view.ndim, 1)
            self.assertEquals(view.shape, (count,))
            self.assertEquals(view.format, fmt)
            self.assertEquals(view.obj, Vec(*test_values[:count]))

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

        vec_info = get_vec_info()

        test_values = [9, 8, 7, 6]

        for Vec, size, count, fmt in vec_info:
            if fmt == "e":
                # GfHalf not supported with array.array
                continue

            test_array = array.array(fmt, test_values[:count])
            test_vec = Vec(test_array)

            self.assertEquals(test_vec, Vec(*test_values[:count]))

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

        import usdrt

        vec_info = get_vec_info()

        equivalent_types = {
            usdrt.Gf.Vec2d: pxr.Gf.Vec2d,
            usdrt.Gf.Vec2f: pxr.Gf.Vec2f,
            usdrt.Gf.Vec2h: pxr.Gf.Vec2h,
            usdrt.Gf.Vec2i: pxr.Gf.Vec2i,
            usdrt.Gf.Vec3d: pxr.Gf.Vec3d,
            usdrt.Gf.Vec3f: pxr.Gf.Vec3f,
            usdrt.Gf.Vec3h: pxr.Gf.Vec3h,
            usdrt.Gf.Vec3i: pxr.Gf.Vec3i,
            usdrt.Gf.Vec4d: pxr.Gf.Vec4d,
            usdrt.Gf.Vec4f: pxr.Gf.Vec4f,
            usdrt.Gf.Vec4h: pxr.Gf.Vec4h,
            usdrt.Gf.Vec4i: pxr.Gf.Vec4i,
        }

        test_values = [9, 8, 7, 6]

        for Vec, size, count, fmt in vec_info:

            test_pxr_type = equivalent_types[Vec](test_values[:count])
            # Note that in all cases pybind choses the tuple initializer and
            # not the buffer initializer. Oh well.
            test_vec = Vec(test_pxr_type)

            self.assertEquals(test_vec, Vec(*test_values[:count]))

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

        vec_info = get_vec_info()

        equivalent_types = get_vec_numpy_types()

        test_values = [9, 8, 7, 6]

        for Vec, size, count, fmt in vec_info:

            test_numpy_type = numpy.array(test_values[:count], dtype=equivalent_types[Vec])
            # Note that in all cases except half pybind choses the tuple initializer and
            # not the buffer initializer. Oh well.
            test_vec = Vec(test_numpy_type)

            self.assertEquals(test_vec, Vec(*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 array

        vec_info = get_vec_info()

        test_values = [9, 8, 7, 6]

        for Vec, size, count, fmt in vec_info:
            if fmt in ["e"]:
                # GfHalf is tricky because it will cast from any
                # numeric type, which pybind too-helpfully interprets
                # into a tuple of GfHalf
                continue

            # Here, pybind will convert arrays of similar types to tuples
            # which is fine I guess but makes testing harder to validate
            # type correctness with buffer support
            test_array = array.array("d" if fmt == "i" else "h", test_values[:count])

            with self.assertRaises(ValueError):
                test_vec = Vec(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

        vec_info = get_vec_info()

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

        for Vec, size, count, fmt in vec_info:
            if fmt in ["e"]:
                # GfHalf not supported with array.array
                continue

            # Here, pybind will convert arrays of similar types to tuples
            # which is fine I guess but makes testing harder to validate
            # type correctness with buffer support
            test_array = array.array(fmt, test_values)

            with self.assertRaises(ValueError):
                test_vec = Vec(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

        vec_info = get_vec_info()
        equivalent_types = get_vec_numpy_types()

        test_values = [9, 8, 7, 6]

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

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


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

    @tc_logger
    def test_copy(self):
        """Test __copy__"""
        vec_info = get_vec_info()

        for Vec, size, count, fmt in vec_info:
            x = Vec()
            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__"""
        vec_info = get_vec_info()

        for Vec, size, count, fmt in vec_info:
            x = Vec()
            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)
