__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 os
import pathlib
import platform
import sys
import time
import unittest

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
        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 = "."


class TestUsdAttribute(TestClass):
    @tc_logger
    def test_get(self):
        from usdrt import Gf, Sdf, Usd, UsdGeom, UsdLux

        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(UsdGeom.Tokens.doubleSided)
        self.assertTrue(attr)

        result = False
        result = attr.Get()
        self.assertTrue(result)
        self.assertTrue(isinstance(result, bool))

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        self.assertTrue(prim)

        attr = prim.GetAttribute(UsdLux.Tokens.inputsIntensity)
        self.assertTrue(attr)

        result = 5.0
        result = attr.Get()
        self.assertEquals(result, 150.0)
        self.assertTrue(isinstance(result, float))

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

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        attr = prim.GetAttribute(UsdLux.Tokens.inputsColor)
        result = attr.Get()
        self.assertTrue(isinstance(result, Gf.Vec3f))
        self.assertEquals(result, Gf.Vec3f(1, 1, 1))

        attr = prim.GetAttribute(UsdGeom.Tokens.visibility)
        result = attr.Get()
        self.assertEquals(result, UsdGeom.Tokens.inherited)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cube_5/subset/BasicShapeMaterial/BasicShapeMaterial"))
        attr = prim.GetAttribute("path")
        result = attr.Get()
        self.assertEquals(result, "/Engine/BasicShapes/BasicShapeMaterial.BasicShapeMaterial")

        # invalid attr should not cause crash
        attr = prim.GetAttribute("invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.Get())

    @tc_logger
    def test_set(self):
        from usdrt import Gf, Sdf, Usd, UsdGeom, UsdLux, Vt

        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(UsdGeom.Tokens.doubleSided)
        self.assertTrue(attr)

        result = False
        result = attr.Get()
        self.assertTrue(result)
        self.assertTrue(isinstance(result, bool))

        attr.Set(False)
        result = attr.Get()
        self.assertFalse(result)
        self.assertTrue(isinstance(result, bool))

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        self.assertTrue(prim)

        attr = prim.GetAttribute(UsdLux.Tokens.inputsIntensity)
        self.assertTrue(attr)

        result = 5.0
        result = attr.Get()
        self.assertEquals(result, 150.0)
        self.assertTrue(isinstance(result, float))

        attr.Set(15.5)
        result = attr.Get()
        self.assertEquals(result, 15.5)
        self.assertTrue(isinstance(result, float))

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))
        self.assertTrue(prim)
        attr = prim.GetAttribute(UsdGeom.Tokens.faceVertexIndices)
        self.assertTrue(attr)
        newValues = Vt.IntArray([9, 8, 7, 6])
        self.assertTrue(attr.Set(newValues))

        result = attr.Get()
        self.assertEquals(len(result), 4)
        self.assertTrue(result.IsFabricData())
        for i in range(4):
            self.assertEquals(result[i], 9 - i)

        newValues = Vt.IntArray([0, 1, 2, 3, 4, 5, 6])
        self.assertTrue(attr.Set(newValues))
        result = attr.Get()
        self.assertEquals(len(result), 7)
        self.assertTrue(result.IsFabricData())
        for i in range(7):
            self.assertEquals(result[i], i)

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        attr = prim.GetAttribute(UsdLux.Tokens.inputsColor)
        newColor = Gf.Vec3f(0, 0.5, 0.5)
        self.assertTrue(attr.Set(newColor))
        vecCheck = attr.Get()
        self.assertTrue(isinstance(vecCheck, Gf.Vec3f))
        self.assertEquals(vecCheck, Gf.Vec3f(0, 0.5, 0.5))

        # setting with an invalid type for the attr should fail
        self.assertFalse(attr.Set("stringvalue"))
        result = attr.Get()
        self.assertEquals(result, Gf.Vec3f(0, 0.5, 0.5))

        attr = prim.GetAttribute(UsdGeom.Tokens.visibility)
        result = attr.Get()
        self.assertEquals(result, UsdGeom.Tokens.inherited)
        self.assertTrue(attr.Set(UsdGeom.Tokens.invisible))
        result = attr.Get()
        self.assertEquals(result, UsdGeom.Tokens.invisible)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cube_5/subset/BasicShapeMaterial/BasicShapeMaterial"))
        attr = prim.GetAttribute("path")
        result = attr.Get()
        self.assertEquals(result, "/Engine/BasicShapes/BasicShapeMaterial.BasicShapeMaterial")
        newString = "/Engine/BasicShapes/BasicShapeMaterial.BasicShapeMaterial100"
        self.assertTrue(attr.Set(newString))
        result = attr.Get()
        self.assertEquals(result, "/Engine/BasicShapes/BasicShapeMaterial.BasicShapeMaterial100")

        # invalid attr should not cause crash
        attr = prim.GetAttribute("invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.Set("stringvalue"))

    @tc_logger
    def test_has_value(self):
        from usdrt import Sdf, Usd, UsdGeom

        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(UsdGeom.Tokens.doubleSided)
        self.assertTrue(attr.HasValue())

        attr2 = prim.GetAttribute(UsdGeom.Tokens.purpose)
        self.assertFalse(attr2.HasValue())

        # invalid attr should not cause crash
        attr = prim.GetAttribute("invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.HasValue())

    @tc_logger
    def test_name(self):
        from usdrt import Sdf, Usd, UsdGeom

        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(UsdGeom.Tokens.primvarsDisplayColor)
        self.assertTrue(attr)
        self.assertTrue(attr.HasValue())
        self.assertEquals(attr.GetName(), UsdGeom.Tokens.primvarsDisplayColor)

        self.assertEquals(attr.GetBaseName(), "displayColor")
        self.assertEquals(attr.GetNamespace(), "primvars")

        parts = attr.SplitName()
        self.assertEquals(len(parts), 2)
        self.assertEquals(parts[0], "primvars")
        self.assertEquals(parts[1], "displayColor")

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

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

        prim = stage.DefinePrim(Sdf.Path("/test"), "Xform")
        self.assertTrue(prim)

        # These are the types that I know are missing support, either
        # because Fabric doesn't support them (string array), or USDRT hasn't implemented
        # support yet (assets), or there will never be a Python rep (PrimTagType)
        KNOWN_INVALID_TYPES = [
            "StringArray",
        ]

        KNOWN_UNSUPPORTED_TYPES = [
            "Frame4d",
            "Frame4dArray",
        ]

        NO_DATA_TAG_TYPES = [
            "PrimTypeTag",
            "AppliedSchemaTypeTag",
            "AncestorPrimTypeTag",
            "Tag",
        ]

        # These are the types that should be supported and are not for some
        # reason - OM-53013
        FIXME_UNSUPPORTED_TYPES = [
            "Matrix2d",
            "Matrix2dArray",
            "TimeCode",
            "TimeCodeArray",
        ]

        for type_name in dir(Sdf.ValueTypeNames):
            if type_name.startswith("__"):
                continue

            if type_name in KNOWN_INVALID_TYPES:
                continue

            value_type_name = getattr(Sdf.ValueTypeNames, type_name)
            if not isinstance(value_type_name, Sdf.ValueTypeName):
                continue

            attr = prim.CreateAttribute(type_name, value_type_name, True)
            self.assertTrue(attr)

            result = attr.Get()
            if result is None:
                # Python bindings warn about unsupported type using C++ name,
                # add a note to clarify using SdfValueTypeName
                self.assertTrue(
                    type_name in KNOWN_UNSUPPORTED_TYPES
                    or type_name in FIXME_UNSUPPORTED_TYPES
                    or type_name in NO_DATA_TAG_TYPES
                )
                if type_name in FIXME_UNSUPPORTED_TYPES:
                    reason = "this will be fixed eventually"
                elif type_name in KNOWN_UNSUPPORTED_TYPES:
                    reason = "this is intentional"
                elif type_name in NO_DATA_TAG_TYPES:
                    reason = "expected no data attribute"
                print(f"i.e. SdfValueTypeName.{type_name} - {reason}")

    @tc_logger
    def test_get_type(self):
        from usdrt import Sdf, Usd, UsdGeom

        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(UsdGeom.Tokens.faceVertexCounts)
        self.assertTrue(attr)
        self.assertEqual(attr.GetTypeName(), Sdf.ValueTypeNames.IntArray)

        attr = prim.GetAttribute(UsdGeom.Tokens.normals)
        self.assertTrue(attr)
        self.assertEqual(attr.GetTypeName(), Sdf.ValueTypeNames.Normal3fArray)

        # invalid attr should not cause crash
        attr = prim.GetAttribute("invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.GetTypeName())

    @tc_logger
    def test_get_type(self):
        from usdrt import Sdf, Usd, UsdGeom

        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(UsdGeom.Tokens.faceVertexCounts)
        self.assertTrue(attr)

        expected = "Attribute(</Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back.faceVertexCounts>)"
        self.assertEqual(repr(attr), expected)

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

        from usdrt import Sdf, Usd, UsdGeom, 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.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(UsdGeom.Tokens.points)
        self.assertTrue(attr)

        self.assertTrue(attr.IsCpuDataValid())
        self.assertFalse(attr.IsGpuDataValid())

        # Force sync to GPU
        attr.SyncDataToGpu()
        self.assertTrue(attr.IsCpuDataValid())
        self.assertTrue(attr.IsGpuDataValid())

        # Write a new GPU value
        points = attr.Get()
        wp_points = wp.array(points, device="cpu")

        with wp.ScopedCudaGuard():

            wpgpu_points = wp_points.to("cuda")

            vt_gpu = Vt.Vec3fArray(wpgpu_points)
            attr.Set(vt_gpu)

            wp.synchronize()

        self.assertFalse(attr.IsCpuDataValid())
        self.assertTrue(attr.IsGpuDataValid())

        # Force sync to CPU
        attr.SyncDataToCpu()
        self.assertTrue(attr.IsCpuDataValid())
        self.assertTrue(attr.IsGpuDataValid())

        # Invalidate data on CPU
        attr.InvalidateCpuData()
        self.assertFalse(attr.IsCpuDataValid())
        self.assertTrue(attr.IsGpuDataValid())

        # Invalidate data on GPU
        attr.InvalidateGpuData()
        self.assertTrue(attr.IsCpuDataValid())
        self.assertFalse(attr.IsGpuDataValid())

    # Test Connections
    @tc_logger
    def test_has_authored_connections(self):
        from usdrt import Sdf, Usd

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

        attr = stage.GetAttributeAtPath(Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:mdl:displacement"))
        self.assertTrue(attr)

        self.assertTrue(attr.HasAuthoredConnections())

        attr = stage.GetAttributeAtPath(Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:surface"))
        self.assertTrue(attr)
        self.assertFalse(attr.HasAuthoredConnections())

        attr = stage.GetAttributeAtPath(
            Sdf.Path("/Cornell_Box/Root/Looks/Green1SG/Green_Side.info:implementationSource")
        )
        self.assertTrue(attr)
        self.assertFalse(attr.HasAuthoredConnections())

        # invalid attr should not cause crash
        attr = stage.GetAttributeAtPath("/Cornell_Box.invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.HasAuthoredConnections())

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

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

        attr = stage.GetAttributeAtPath(Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:mdl:displacement"))
        self.assertTrue(attr)

        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], Sdf.Path("/Cornell_Box/Root/Looks/Green1SG/Green_Side.outputs:out"))

        # invalid attr should not cause crash
        attr = stage.GetAttributeAtPath("/Cornell_Box.invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.GetConnections())

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

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)
        path = Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:mdl:displacement")
        attr = stage.GetAttributeAtPath(path)
        self.assertTrue(attr)

        new_connection = Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:surface")
        self.assertTrue(attr.AddConnection(new_connection))

        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], new_connection)

        # add a connection to a new attribute
        proxy_prim_connection = Sdf.Path("/Cube_5")
        prim = stage.GetPrimAtPath("/Cornell_Box")
        self.assertTrue(prim)
        attr = prim.CreateAttribute("testAttr", Sdf.ValueTypeNames.Token, True)
        self.assertTrue(attr)
        self.assertTrue(attr.AddConnection(proxy_prim_connection))
        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], Sdf.Path("/Cube_5"))

        # invalid attr should not cause crash
        attr = stage.GetAttributeAtPath("/Cornell_Box.invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.AddConnection(Sdf.Path("/Cube_5")))

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

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

        attr = stage.GetAttributeAtPath(Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:mdl:displacement"))
        self.assertTrue(attr)

        # It's valid to try to remove a connection that is not
        # in the list of connections
        connection_not_in_list = Sdf.Path("/Foo/Bar")
        self.assertTrue(attr.RemoveConnection(connection_not_in_list))

        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], Sdf.Path("/Cornell_Box/Root/Looks/Green1SG/Green_Side.outputs:out"))

        # Removing a valid connection will return empty connection lists with GetConnections
        connection_in_list = Sdf.Path("/Cornell_Box/Root/Looks/Green1SG/Green_Side.outputs:out")
        self.assertTrue(attr.RemoveConnection(connection_in_list))

        connections = attr.GetConnections()
        self.assertEqual(len(connections), 0)

        # multiple connections not supported yet
        # test removing if there are multiple in the list
        # multiple_connections = [Sdf.Path("/Foo"), Sdf.Path("/Bar")]
        # self.assertTrue(attr.SetConnections(multiple_connections))
        # self.assertTrue(attr.RemoveConnection(Sdf.Path("/Bar")))
        # connections = attr.GetConnections()
        # self.assertEqual(len(connections), 1)
        # self.assertEqual(connections[0], Sdf.Path("/Foo"))

        # invalid attr should not cause crash
        attr = stage.GetAttributeAtPath("/Cornell_Box.invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.RemoveConnection(Sdf.Path("/Cube_5")))

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

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

        attr = stage.GetAttributeAtPath(Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:mdl:displacement"))
        self.assertTrue(attr)

        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], Sdf.Path("/Cornell_Box/Root/Looks/Green1SG/Green_Side.outputs:out"))

        self.assertTrue(attr.ClearConnections())
        connections = attr.GetConnections()
        self.assertEqual(len(connections), 0)

        # invalid attr should not cause crash
        attr = stage.GetAttributeAtPath("/Cornell_Box.invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.ClearConnections())

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

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

        attr = stage.GetAttributeAtPath(Sdf.Path("/Cornell_Box/Root/Looks/Green1SG.outputs:mdl:displacement"))
        self.assertTrue(attr)

        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], Sdf.Path("/Cornell_Box/Root/Looks/Green1SG/Green_Side.outputs:out"))

        new_connections = [Sdf.Path("/Cornell_Box/Root/Looks/Red1SG/Red_Side.outputs:out")]
        self.assertTrue(attr.SetConnections(new_connections))
        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(connections[0], Sdf.Path("/Cornell_Box/Root/Looks/Red1SG/Red_Side.outputs:out"))

        too_many_connections = [Sdf.Path("/Foo"), Sdf.Path("/Bar")]
        # do nothing for now, this isnt supported
        self.assertFalse(attr.SetConnections(too_many_connections))
        connections = attr.GetConnections()
        self.assertEqual(len(connections), 1)
        # same as previous connection
        self.assertEqual(connections[0], Sdf.Path("/Cornell_Box/Root/Looks/Red1SG/Red_Side.outputs:out"))

        # invalid attr should not cause crash
        attr = stage.GetAttributeAtPath("/Cornell_Box.invalid")
        self.assertFalse(attr)
        self.assertFalse(attr.SetConnections(new_connections))

