__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 ctypes
import pathlib
import platform
import weakref

from . import tc_logger


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()

        if platform.processor() == "aarch64":
            # warp not supported on aarch64 yet for our testing
            return

        import warp as wp

        wp.init()
    except ImportError:
        # not needed in Kit /
        # warp not available in default kit install
        # not needed in Kit
        pass


try:
    import omni.kit.test

    TestClass = omni.kit.test.AsyncTestCase
    TEST_DIR = str(pathlib.Path(__file__).parent.parent.parent.parent)
except ImportError:
    from . import TestRtScenegraph

    TestClass = TestRtScenegraph
    TEST_DIR = "."


def warp_array_from_cuda_array_interface(a):
    import warp as wp

    cai = a.__cuda_array_interface__
    return wp.types.array(
        dtype=wp.vec3,
        length=cai["shape"][0],
        capacity=cai["shape"][0] * wp.types.type_size_in_bytes(wp.vec3),
        ptr=cai["data"][0],
        device="cuda",
        owner=False,
        requires_grad=False,
    )


class TestVtArray(TestClass):
    @tc_logger
    def test_init(self):
        from usdrt import Vt

        empty = Vt.IntArray()
        self.assertEquals(len(empty), 0)

        allocated = Vt.IntArray(10)
        self.assertEquals(len(allocated), 10)

        from_list = Vt.IntArray([0, 1, 2, 3])
        self.assertEquals(len(from_list), 4)

    @tc_logger
    def test_getset(self):
        from usdrt import Vt

        test_list = Vt.IntArray([0, 1, 2, 3])
        self.assertEquals(len(test_list), 4)

        for i in range(4):
            self.assertEquals(test_list[i], i)

        test_list[3] = 100
        self.assertEquals(test_list[3], 100)

    @tc_logger
    def test_iterator(self):
        from usdrt import Vt

        test_list = Vt.IntArray([0, 1, 2, 3])
        i = 0
        for item in test_list:
            self.assertEquals(i, item)
            i += 1

    @tc_logger
    def test_is_fabric(self):
        from usdrt import Sdf, Usd, Vt

        test_list = Vt.IntArray([0, 1, 2, 3])
        self.assertFalse(test_list.IsPythonData())

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))
        self.assertTrue(prim)
        attr = prim.GetAttribute("faceVertexIndices")
        self.assertTrue(attr)
        result = attr.Get()
        self.assertEquals(result[0], 1)
        self.assertTrue(result.IsFabricData())
        result[0] = 9
        self.assertTrue(result.IsFabricData())

        result2 = attr.Get()
        self.assertTrue(result2.IsFabricData())
        self.assertEquals(result2[0], 9)

        result2.DetachFromSource()
        result2[0] = 15
        result3 = attr.Get()
        self.assertEquals(result3[0], 9)

    def _getNumpyPoints(self):
        import numpy as np

        points = np.zeros(shape=(1024, 3), dtype=np.float32)
        points[0][0] = 999.0
        points[0][1] = 998.0
        points[0][2] = 997.0
        points[1][0] = 1.0
        points[1][1] = 2.0
        points[1][2] = 3.0
        points[10][0] = 10.0
        points[10][1] = 20.0
        points[10][2] = 30.0
        return points

    @tc_logger
    def test_array_interface(self):
        import numpy as np

        from usdrt import Sdf, Usd, Vt

        stage = Usd.Stage.CreateInMemory("test.usda")
        self.assertTrue(stage)

        prim = stage.DefinePrim(Sdf.Path("/prim"))
        attr = prim.CreateAttribute("points", Sdf.ValueTypeNames.Point3fArray, True)
        self.assertTrue(attr)
        points = self._getNumpyPoints()
        pointsRef = weakref.ref(points)
        pointsId = id(points)

        # Create a VtArray from a numpy array using the buffer protocol
        vtPoints = Vt.Vec3fArray(points)

        # Ensure the VtArray incremented the refcount on points
        self.assertEqual(ctypes.c_long.from_address(pointsId).value, 2)

        # Delete the numpy reference to check that the VtArray kept it alive
        del points
        # Ensure VtArray released the refcount on points
        self.assertEqual(ctypes.c_long.from_address(pointsId).value, 1)

        # Set the Fabric attribute which will make a copy from numpy to Fabric
        self.assertTrue(attr.Set(vtPoints))

        del vtPoints
        # Ensure VtArray released the refcount on points
        self.assertIsNone(pointsRef())

        # Retrieve a new Fabric VtArray from usdrt
        fabricPoints = attr.Get()
        self.assertTrue(fabricPoints.IsFabricData())

        # Create a new attribute which would invalidate any existing Fabric pointer
        attr2 = prim.CreateAttribute("badattr", Sdf.ValueTypeNames.Float, True)
        self.assertTrue(attr2.Set(3.5))

        newPoints = np.array(fabricPoints)
        # Delete the fabric VtArray to ensure the data was really copied into numpy
        del fabricPoints

        # Now compare the round-trip through usdrt was a success
        points = self._getNumpyPoints()
        for i in range(1024):
            for j in range(3):
                self.assertEqual(
                    points[i][j], newPoints[i][j], msg=f"points[{i}][{j}]: {points[i][j]} != {newPoints[i][j]}"
                )

    @tc_logger
    def test_cuda_array_interface(self):
        if platform.processor() == "aarch64":
            return

        import numpy as np

        from usdrt import Sdf, Usd, Vt

        try:
            import warp as wp
        except ImportError:
            if TestClass is omni.kit.test.AsyncTestCase:
                # warp not available for testing in kit at this time
                return
            raise

        stage = Usd.Stage.CreateInMemory("test.usda")
        self.assertTrue(stage)

        prim = stage.DefinePrim(Sdf.Path("/prim"))
        attr = prim.CreateAttribute("points", Sdf.ValueTypeNames.Point3fArray, True)
        self.assertTrue(attr)
        idattr = prim.CreateAttribute("faceVertexIndices", Sdf.ValueTypeNames.IntArray, True)
        self.assertTrue(idattr)
        tag = prim.CreateAttribute("Deformable", Sdf.ValueTypeNames.PrimTypeTag, True)
        self.assertTrue(tag)

        # First create a warp array from numpy
        points = self._getNumpyPoints()
        warpPoints = wp.array(points, device="cpu")

        with wp.ScopedCudaGuard():
            gpuWarpPoints = warpPoints.to("cuda")
            warpId = id(gpuWarpPoints)
            warpRef = weakref.ref(gpuWarpPoints)

            # Create a VtArray from a warp array using the __cuda_array_interface__
            self.assertEqual(ctypes.c_long.from_address(warpId).value, 1)
            vtPoints = Vt.Vec3fArray(gpuWarpPoints)
            self.assertEqual(ctypes.c_long.from_address(warpId).value, 2)

            # Delete the warp reference to check that the VtArray kept it alive
            del gpuWarpPoints
            self.assertEqual(ctypes.c_long.from_address(warpId).value, 1)

            # Set the Fabric attribute which will make a CUDA copy from warp to Fabric
            self.assertTrue(attr.Set(vtPoints))

            # OM-84460 - IntArray supported for 1d warp arrays
            warp_ids = wp.zeros(shape=1024, dtype=int, device="cuda")
            vt_ids = Vt.IntArray(warp_ids)
            self.assertTrue(idattr.Set(vt_ids))

            wp.synchronize()

            # Check the VtArray destruction released the warp array
            del vtPoints
            self.assertIsNone(warpRef())

            # Retrieve a new Fabric VtArray from usdrt
            newVtPoints = attr.Get()
            self.assertTrue(newVtPoints.HasFabricGpuData())

            # Create a new attribute which would invalidate any existing Fabric pointer
            attr2 = prim.CreateAttribute("badattr", Sdf.ValueTypeNames.Float, True)
            attr2.Set(3.5)

            newWarpPoints = warp_array_from_cuda_array_interface(newVtPoints)

            # Delete the fabric VtArray to ensure the data was really copied into warp
            del newVtPoints

            newPoints = newWarpPoints.numpy()

        # Now compare the round-trip through usdrt was a success
        self.assertEqual(points.shape, newPoints.shape)
        for i in range(1024):
            for j in range(3):
                self.assertEqual(
                    points[i][j], newPoints[i][j], msg=f"points[{i}][{j}]: {points[i][j]} != {newPoints[i][j]}"
                )

    @tc_logger
    def test_repr(self):
        from usdrt import Vt

        test_array = Vt.IntArray(4)

        result = "Vt.IntArray(4, ("

        for i in range(4):
            test_array[i] = i
            result += str(i)
            if i < 3:
                result += ", "
            else:
                result += ")"

        result += ")"

        self.assertEqual(result, repr(test_array))

    @tc_logger
    def test_str(self):
        from usdrt import Vt
        from pxr import Vt as PxrVt

        test_array = Vt.IntArray(range(4))

        expected = str(PxrVt.IntArray(range(4)))

        self.assertEqual(str(test_array), expected)

    @tc_logger
    def test_assetarray(self):
        from usdrt import Sdf, Vt

        test_array = Vt.AssetArray(4)

        for i in range(4):
            test_array[i] = Sdf.AssetPath(f"hello{i}", f"hello{i}.txt")

        self.assertEqual(len(test_array), 4)

        # This should not cause a crash
        del test_array

    @tc_logger
    def test_arraybuffer_types(self):
        import numpy as np

        from usdrt import Vt

        self.assertTrue(Vt.BoolArray(np.zeros(shape=(10, 1), dtype=np.bool_)).IsPythonData())
        self.assertTrue(Vt.CharArray(np.zeros(shape=(10, 1), dtype=np.byte)).IsPythonData())
        self.assertTrue(Vt.UCharArray(np.zeros(shape=(10, 1), dtype=np.ubyte)).IsPythonData())
        self.assertTrue(Vt.DoubleArray(np.zeros(shape=(10, 1), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.FloatArray(np.zeros(shape=(10, 1), dtype=np.float32)).IsPythonData())
        self.assertTrue(Vt.HalfArray(np.zeros(shape=(10, 1), dtype=np.float16)).IsPythonData())
        self.assertTrue(Vt.IntArray(np.zeros(shape=(10, 1), dtype=np.intc)).IsPythonData())
        self.assertTrue(Vt.Int64Array(np.zeros(shape=(10, 1), dtype=np.longlong)).IsPythonData())
        self.assertTrue(Vt.ShortArray(np.zeros(shape=(10, 1), dtype=np.short)).IsPythonData())
        self.assertTrue(Vt.UInt64Array(np.zeros(shape=(10, 1), dtype=np.ulonglong)).IsPythonData())
        self.assertTrue(Vt.UIntArray(np.zeros(shape=(10, 1), dtype=np.uintc)).IsPythonData())
        self.assertTrue(Vt.UShortArray(np.zeros(shape=(10, 1), dtype=np.ushort)).IsPythonData())

        self.assertTrue(Vt.Matrix3dArray(np.zeros(shape=(10, 9), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.Matrix3fArray(np.zeros(shape=(10, 9), dtype=np.float32)).IsPythonData())
        self.assertTrue(Vt.Matrix4dArray(np.zeros(shape=(10, 16), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.Matrix4fArray(np.zeros(shape=(10, 16), dtype=np.float32)).IsPythonData())

        self.assertTrue(Vt.QuatdArray(np.zeros(shape=(10, 4), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.QuatfArray(np.zeros(shape=(10, 4), dtype=np.float32)).IsPythonData())
        self.assertTrue(Vt.QuathArray(np.zeros(shape=(10, 4), dtype=np.float16)).IsPythonData())

        self.assertTrue(Vt.Vec2dArray(np.zeros(shape=(10, 2), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.Vec3dArray(np.zeros(shape=(10, 3), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.Vec4dArray(np.zeros(shape=(10, 4), dtype=np.float64)).IsPythonData())
        self.assertTrue(Vt.Vec2fArray(np.zeros(shape=(10, 2), dtype=np.float32)).IsPythonData())
        self.assertTrue(Vt.Vec3fArray(np.zeros(shape=(10, 3), dtype=np.float32)).IsPythonData())
        self.assertTrue(Vt.Vec4fArray(np.zeros(shape=(10, 4), dtype=np.float32)).IsPythonData())
        self.assertTrue(Vt.Vec2hArray(np.zeros(shape=(10, 2), dtype=np.float16)).IsPythonData())
        self.assertTrue(Vt.Vec3hArray(np.zeros(shape=(10, 3), dtype=np.float16)).IsPythonData())
        self.assertTrue(Vt.Vec4hArray(np.zeros(shape=(10, 4), dtype=np.float16)).IsPythonData())
        self.assertTrue(Vt.Vec2iArray(np.zeros(shape=(10, 2), dtype=np.intc)).IsPythonData())
        self.assertTrue(Vt.Vec3iArray(np.zeros(shape=(10, 3), dtype=np.intc)).IsPythonData())
        self.assertTrue(Vt.Vec4iArray(np.zeros(shape=(10, 4), dtype=np.intc)).IsPythonData())

    @tc_logger
    def test_implicit_conversion_bindings(self):
        from usdrt import Gf, Sdf, Usd, Vt

        stage = Usd.Stage.CreateInMemory("test.usda")
        self.assertTrue(stage)

        prim = stage.DefinePrim(Sdf.Path("/prim"))

        # IntArray
        intarray = prim.CreateAttribute("intarray", Sdf.ValueTypeNames.IntArray, True)
        self.assertTrue(intarray)

        intarray.Set(Vt.IntArray([0, 1, 2]))
        intlist = list(prim.GetAttribute("intarray").Get())
        self.assertEqual(len(intlist), 3)
        intlist.extend([3])
        self.assertEqual(len(intlist), 4)
        self.assertTrue(prim.GetAttribute("intarray").Set(intlist))

        updated_list = prim.GetAttribute("intarray").Get()
        self.assertEqual(len(updated_list), 4)
        self.assertEqual(list(updated_list), [0, 1, 2, 3])

        # Point3fArray
        points = prim.CreateAttribute("points", Sdf.ValueTypeNames.Point3fArray, True)
        self.assertTrue(points)

        points.Set(Vt.Vec3fArray([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]))

        vertices = list(prim.GetAttribute("points").Get())
        self.assertEqual(len(vertices), 3)
        vertices.extend([[3.0, 3.0, 3.0]])

        prim.GetAttribute("points").Set(vertices)

        updated_vertices = prim.GetAttribute("points").Get()
        self.assertEqual(len(updated_vertices), 4)

