__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()
    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 TestRtChangeTracker(TestClass):
    @tc_logger
    def test_ctor(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        tracker = Rt.ChangeTracker(stage)
        self.assertTrue(tracker)

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_track_attribute(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        tracker = Rt.ChangeTracker(stage)
        self.assertTrue(tracker)

        tracker.TrackAttribute("inputs:color")
        self.assertTrue(tracker.IsTrackingAttribute("inputs:color"))
        self.assertFalse(tracker.HasChanges())

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        attr = prim.GetAttribute("inputs:color")

        newColor = Gf.Vec3f(0, 0.5, 0.5)
        attr.Set(newColor, Usd.TimeCode(0.0))
        self.assertTrue(tracker.HasChanges())

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_stop_tracking_attribute(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        attr = prim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        self.assertTrue(tracker)
        tracker.TrackAttribute("inputs:color")
        self.assertTrue(tracker.IsTrackingAttribute("inputs:color"))
        self.assertFalse(tracker.HasChanges())

        tracker.StopTrackingAttribute("inputs:color")
        self.assertFalse(tracker.IsTrackingAttribute("inputs:color"))

        newColor = Gf.Vec3f(0, 0.6, 0.6)
        attr.Set(newColor, Usd.TimeCode(0.0))
        self.assertFalse(tracker.HasChanges())

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_pause_and_resume_tracking(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        attr = prim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        self.assertTrue(tracker)
        tracker.TrackAttribute("inputs:color")
        self.assertTrue(tracker.IsTrackingAttribute("inputs:color"))

        tracker.PauseTracking()
        newColorA = Gf.Vec3f(0, 0.7, 0.7)
        attr.Set(newColorA, Usd.TimeCode(0.0))
        self.assertFalse(tracker.HasChanges())

        tracker.ResumeTracking()
        newColorB = Gf.Vec3f(0, 0.8, 0.8)
        attr.Set(newColorB, Usd.TimeCode(0.0))
        self.assertTrue(tracker.HasChanges())

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_is_change_tracking_paused(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        tracker = Rt.ChangeTracker(stage)
        self.assertTrue(tracker)
        tracker.TrackAttribute("inputs:color")

        self.assertFalse(tracker.IsChangeTrackingPaused())
        tracker.PauseTracking()
        self.assertTrue(tracker.IsChangeTrackingPaused())
        tracker.ResumeTracking()
        self.assertFalse(tracker.IsChangeTrackingPaused())

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_has_and_clear_changes(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        prim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        attr = prim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        tracker.TrackAttribute("inputs:color")
        self.assertFalse(tracker.HasChanges())

        newColor = Gf.Vec3f(0, 0.9, 0.9)
        attr.Set(newColor, Usd.TimeCode(0.0))
        self.assertTrue(tracker.HasChanges())

        tracker.ClearChanges()
        self.assertFalse(tracker.HasChanges())

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_get_all_changed_prims(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        lightPrim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        colorAttr = lightPrim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        tracker.TrackAttribute("inputs:color")
        tracker.TrackAttribute("visibility")

        # No changes
        changedPrims = tracker.GetAllChangedPrims()
        self.assertEqual(len(changedPrims), 0)

        # Change on one prim
        newColor = Gf.Vec3f(0, 0.4, 0.4)
        colorAttr.Set(newColor, Usd.TimeCode(0.0))
        changedPrims = tracker.GetAllChangedPrims()
        self.assertEqual(len(changedPrims), 1)
        self.assertEqual(changedPrims[0], Sdf.Path("/DistantLight"))

        # Change on two prims
        cubePrim = stage.GetPrimAtPath(Sdf.Path("/Cube_5"))
        visAttr = cubePrim.GetAttribute("visibility")
        visAttr.Set("invisible", Usd.TimeCode(0.0))
        changedPrims = tracker.GetAllChangedPrims()
        self.assertEqual(len(changedPrims), 2)

        # Change on a prim that already has a change
        lightVisAttr = lightPrim.GetAttribute("visibility")
        lightVisAttr.Set("invisible", Usd.TimeCode(0.0))
        changedPrims = tracker.GetAllChangedPrims()
        self.assertEqual(len(changedPrims), 2)

        # Change an untracked attr on a third prim
        spherePrim = stage.GetPrimAtPath(Sdf.Path("/Sphere_2"))
        displayColor = spherePrim.GetAttribute("primvars:displayColor")
        newDisplayColor = Gf.Vec3f(0, 0.3, 0.3)
        displayColor.Set(newDisplayColor)
        changedPrims = tracker.GetAllChangedPrims()
        self.assertEqual(len(changedPrims), 2)

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_get_all_changed_attributes(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        lightPrim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        colorAttr = lightPrim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        tracker.TrackAttribute("inputs:color")
        tracker.TrackAttribute("visibility")

        # No changes
        changedAttrs = tracker.GetAllChangedAttributes()
        self.assertEqual(len(changedAttrs), 0)

        # Change attr on one prim
        newColor = Gf.Vec3f(0.5, 0.6, 0.7)
        colorAttr.Set(newColor, Usd.TimeCode(0.0))
        changedAttrs = tracker.GetAllChangedAttributes()
        self.assertEqual(len(changedAttrs), 1)
        self.assertEqual(changedAttrs[0], "/DistantLight.inputs:color")

        # Change attr on a second prim
        cubePrim = stage.GetPrimAtPath(Sdf.Path("/Cube_5"))
        cubeVisAttr = cubePrim.GetAttribute("visibility")
        cubeVisAttr.Set("invisible", Usd.TimeCode(0.0))
        changedAttrs = tracker.GetAllChangedAttributes()
        self.assertEqual(len(changedAttrs), 2)

        # Change a second attr on one of the prims
        lightVisAttr = lightPrim.GetAttribute("visibility")
        lightVisAttr.Set("invisible", Usd.TimeCode(0.0))
        changedAttrs = tracker.GetAllChangedAttributes()
        self.assertEqual(len(changedAttrs), 3)

        # Change an untracked attr on a third prim
        spherePrim = stage.GetPrimAtPath(Sdf.Path("/Sphere_2"))
        displayColor = spherePrim.GetAttribute("primvars:displayColor")
        newDisplayColor = Gf.Vec3f(0, 0.3, 0.3)
        displayColor.Set(newDisplayColor, Usd.TimeCode(0.0))
        changedAttrs = tracker.GetAllChangedAttributes()
        self.assertEqual(len(changedAttrs), 3)

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_prim_changed(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        lightPrim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        colorAttr = lightPrim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        tracker.TrackAttribute("inputs:color")
        tracker.TrackAttribute("visibility")

        # UsdPrim version of PrimChanged()
        self.assertFalse(tracker.PrimChanged(lightPrim))
        newColor = Gf.Vec3f(0.6, 0.7, 0.8)
        colorAttr.Set(newColor, Usd.TimeCode(0.0))
        self.assertTrue(tracker.PrimChanged(lightPrim))

        # SdfPath version of PrimChanged()
        cubePrim = stage.GetPrimAtPath(Sdf.Path("/Cube_5"))
        self.assertFalse(tracker.PrimChanged(Sdf.Path("/Cube_5")))
        cubeVisAttr = cubePrim.GetAttribute("visibility")
        cubeVisAttr.Set("invisible", Usd.TimeCode(0.0))
        self.assertTrue(tracker.PrimChanged(Sdf.Path("/Cube_5")))

        # Change an untracked attr on a third prim
        spherePrim = stage.GetPrimAtPath(Sdf.Path("/Sphere_2"))
        displayColor = spherePrim.GetAttribute("primvars:displayColor")
        newDisplayColor = Gf.Vec3f(0, 0.4, 0.4)
        displayColor.Set(newDisplayColor, Usd.TimeCode(0.0))
        self.assertFalse(tracker.PrimChanged(spherePrim))

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_attribute_changed(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        lightPrim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        colorAttr = lightPrim.GetAttribute("inputs:color")

        tracker = Rt.ChangeTracker(stage)
        tracker.TrackAttribute("inputs:color")
        tracker.TrackAttribute("visibility")

        # UsdPrim version of AttributeChanged()
        self.assertFalse(tracker.AttributeChanged(colorAttr))
        newColor = Gf.Vec3f(0.7, 0.8, 0.9)
        colorAttr.Set(newColor, Usd.TimeCode(0.0))
        self.assertTrue(tracker.AttributeChanged(colorAttr))

        # SdfPath version of AttributeChanged()
        cubePrim = stage.GetPrimAtPath(Sdf.Path("/Cube_5"))
        self.assertFalse(tracker.AttributeChanged(Sdf.Path("/Cube_5.visibility")))
        cubeVisAttr = cubePrim.GetAttribute("visibility")
        cubeVisAttr.Set("invisible", Usd.TimeCode(0.0))
        self.assertTrue(tracker.AttributeChanged(Sdf.Path("/Cube_5.visibility")))

        # Change an untracked attr on a third prim
        spherePrim = stage.GetPrimAtPath(Sdf.Path("/Sphere_2"))
        displayColor = spherePrim.GetAttribute("primvars:displayColor")
        newDisplayColor = Gf.Vec3f(0, 0.5, 0.5)
        displayColor.Set(newDisplayColor, Usd.TimeCode(0.0))
        self.assertFalse(tracker.AttributeChanged(displayColor))

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_get_changed_attributes(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)
        lightPrim = stage.GetPrimAtPath(Sdf.Path("/DistantLight"))
        cubePrim = stage.GetPrimAtPath(Sdf.Path("/Cube_5"))
        spherePrim = stage.GetPrimAtPath(Sdf.Path("/Sphere_2"))

        tracker = Rt.ChangeTracker(stage)
        tracker.TrackAttribute("inputs:color")
        tracker.TrackAttribute("visibility")
        tracker.TrackAttribute("subdivisionScheme")

        # Change one attr on a prim
        lightColor = lightPrim.GetAttribute("inputs:color")
        self.assertEqual(len(tracker.GetChangedAttributes(lightPrim)), 0)
        newColor = Gf.Vec3f(0.2, 0.3, 0.4)
        lightColor.Set(newColor, Usd.TimeCode(0.0))
        testResult = tracker.GetChangedAttributes(lightPrim)
        self.assertEqual(len(testResult), 1)
        self.assertEqual(testResult[0], "inputs:color")

        # Change multiple tracked attrs on a prim
        cubeVis = cubePrim.GetAttribute("visibility")
        cubeVis.Set("invisible", Usd.TimeCode(0.0))
        cubeSubdiv = cubePrim.GetAttribute("subdivisionScheme")
        cubeSubdiv.Set("loop", Usd.TimeCode(0.0))
        self.assertEqual(len(tracker.GetChangedAttributes(cubePrim)), 2)

        # Change untracked attr on a prim
        displayColor = spherePrim.GetAttribute("primvars:displayColor")
        newDisplayColor = Gf.Vec3f(0, 0.3, 0.5)
        displayColor.Set(newDisplayColor, Usd.TimeCode(0.0))
        self.assertEqual(len(tracker.GetChangedAttributes(spherePrim)), 0)

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_is_tracking_attribute(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        tracker = Rt.ChangeTracker(stage)
        self.assertFalse(tracker.IsTrackingAttribute("inputs:color"))
        tracker.TrackAttribute("inputs:color")
        self.assertTrue(tracker.IsTrackingAttribute("inputs:color"))

        tracker.TrackAttribute("visibility")
        self.assertTrue(tracker.IsTrackingAttribute("visibility"))
        self.assertTrue(tracker.IsTrackingAttribute("inputs:color"))

        tracker.StopTrackingAttribute("inputs:color")
        self.assertTrue(tracker.IsTrackingAttribute("visibility"))
        self.assertFalse(tracker.IsTrackingAttribute("inputs:color"))
        tracker.StopTrackingAttribute("visibility")
        self.assertFalse(tracker.IsTrackingAttribute("visibility"))

        self.assertFalse(tracker.IsTrackingAttribute("fakeAttr"))

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()

    @tc_logger
    def test_get_tracked_attributes(self):
        from pxr import UsdUtils

        from usdrt import Gf, Rt, Sdf, Usd

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

        tracker = Rt.ChangeTracker(stage)
        self.assertEqual(len(tracker.GetTrackedAttributes()), 0)

        tracker.TrackAttribute("inputs:color")
        self.assertEqual(len(tracker.GetTrackedAttributes()), 1)
        self.assertEqual(tracker.GetTrackedAttributes()[0], "inputs:color")
        tracker.TrackAttribute("visibility")
        self.assertEqual(len(tracker.GetTrackedAttributes()), 2)
        tracker.StopTrackingAttribute("inputs:color")
        self.assertEqual(len(tracker.GetTrackedAttributes()), 1)
        self.assertEqual(tracker.GetTrackedAttributes()[0], "visibility")

        # Really make sure the stage is dead and cache is cleared for test independence
        del stage
        UsdUtils.StageCache.Get().Clear()
