"""Basic tests of the compute graph"""

import unittest

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from pxr import Gf, UsdGeom


# ======================================================================
class TestGatherScatter(ogts.OmniGraphTestCase):
    """Run tests that exercises gather / scatter functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------
    @unittest.skip("Crashes on Linux")
    async def test_simple_gather_scatter(self):
        """
        Basic idea of the test: we have two cubes in the scene.  We gather their translation and velocity
        attributes, bounce them per frame, and then scatter the resulting position back to the cubes
        """
        with og.Settings.temporary(og.USE_SCHEMA_PRIMS_SETTING, True):
            with og.Settings.temporary(og.ALLOW_IMPLICIT_GRAPH_SETTING, True):
                (result, error) = await ogts.load_test_file(
                    "TestBounceGatherScatter.usda", use_caller_subdirectory=True
                )
                self.assertTrue(result, str(error))

                graph = omni.graph.core.get_current_graph()

                box0_node = graph.get_node("/Stage/box_0_0")
                translate_attr = box0_node.get_attribute("xformOp:translate")

                translate = og.Controller.get(translate_attr)
                # ! must cache the value as translate is a reference so will automatically be mutated by the next frame
                t0 = translate[2]

                await omni.kit.app.get_app().next_update_async()
                translate = og.Controller.get(translate_attr)

                diff = abs(translate[2] - t0)
                self.assertGreater(diff, 0.0)

    # ----------------------------------------------------------------------
    async def test_gather_by_path(self):
        """
        Basic idea of the test: we have three cubes in the scene.  We use a node to query their paths by type,
        then we feed this to the gather by path node, which vectorizes the data.  Finally the vectorized data is
        passed to a sample bounce node that does not use bundles.  Instead it relies on the bucketId passed down
        to bounce the cubes.
        """
        with og.Settings.temporary(og.Settings.UPDATE_TO_USD, True):
            (result, error) = await ogts.load_test_file("TestGatherByPath.usda", use_caller_subdirectory=True)
            self.assertTrue(result, str(error))

            usd_context = omni.usd.get_context()
            stage = usd_context.get_stage()
            cube_prim = stage.GetPrimAtPath("/World/Cube_00")
            cube_translate_attr = cube_prim.GetAttribute("xformOp:translate")
            t0 = cube_translate_attr.Get()[2]

            await omni.kit.app.get_app().next_update_async()

            t1 = cube_translate_attr.Get()[2]
            diff = abs(t1 - t0)
            self.assertGreater(diff, 0.0)

    # ----------------------------------------------------------------------
    async def test_gatherprototype_set_gathered_attribute(self):
        """Tests the SetGatheredAttribute node, part of GatherPrototype"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (test_set_node, _, _, _), (_, xform1, xform2, xform_tagged1, _), _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("TestSet", "omni.graph.nodes.SetGatheredAttribute"),
                    ("GatherByPath", "omni.graph.nodes.GatherByPath"),
                ],
                keys.CREATE_PRIMS: [
                    ("Empty", {}),
                    ("Xform1", {"_translate": ("pointd[3]", [0, 0, 0])}),
                    ("Xform2", {"_translate": ("pointd[3]", [0, 0, 0])}),
                    ("XformTagged1", {"foo": ("token", ""), "_translate": ("pointd[3]", [0, 0, 0])}),
                    ("Tagged1", {"foo": ("token", "")}),
                ],
                keys.EXPOSE_PRIMS: [
                    (og.Controller.PrimExposureType.AS_WRITABLE, "/Xform1", "Write1"),
                    (og.Controller.PrimExposureType.AS_WRITABLE, "/Xform2", "Write2"),
                ],
                keys.CONNECT: ("GatherByPath.outputs:gatherId", "TestSet.inputs:gatherId"),
                keys.SET_VALUES: [
                    ("GatherByPath.inputs:primPaths", ["/Xform1", "/Xform2"]),
                    ("GatherByPath.inputs:attributes", "_translate"),
                    ("GatherByPath.inputs:allAttributes", False),
                    ("GatherByPath.inputs:shouldWriteBack", True),
                    ("TestSet.inputs:name", "_translate"),
                    ("TestSet.inputs:value", [[1, 2, 3], [4, 5, 6]], "double[3][]"),
                ],
            },
        )

        await controller.evaluate(graph)
        self.assertEqual(
            Gf.Vec3d(*og.Controller.get(controller.attribute("inputs:value", test_set_node))[0]), Gf.Vec3d(1, 2, 3)
        )
        self.assertEqual(
            Gf.Vec3d(*og.Controller.get(controller.attribute("inputs:value", test_set_node))[1]), Gf.Vec3d(4, 5, 6)
        )
        self.assertEqual(Gf.Vec3d(*xform1.GetAttribute("_translate").Get()), Gf.Vec3d(1, 2, 3))
        self.assertEqual(Gf.Vec3d(*xform2.GetAttribute("_translate").Get()), Gf.Vec3d(4, 5, 6))
        self.assertEqual(Gf.Vec3d(*xform_tagged1.GetAttribute("_translate").Get()), Gf.Vec3d(0, 0, 0))

        # Test write mask
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ("TestSet.inputs:mask", [True, False]),
                    ("TestSet.inputs:value", [[0, 0, 0], [0, 0, 0]], "double[3][]"),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(Gf.Vec3d(*xform1.GetAttribute("_translate").Get()), Gf.Vec3d(0, 0, 0))
        self.assertEqual(Gf.Vec3d(*xform2.GetAttribute("_translate").Get()), Gf.Vec3d(4, 5, 6))

        # Test scaler broadcast
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ("TestSet.inputs:mask", [True, True]),
                    ("TestSet.inputs:value", [1, 2, 3], "double[3]"),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(Gf.Vec3d(*xform1.GetAttribute("_translate").Get()), Gf.Vec3d(1, 2, 3))
        self.assertEqual(Gf.Vec3d(*xform2.GetAttribute("_translate").Get()), Gf.Vec3d(1, 2, 3))

        # Test scaler broadcast + mask
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ("TestSet.inputs:mask", [False, True]),
                    ("TestSet.inputs:value", [1, 1, 1], "double[3]"),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(Gf.Vec3d(*xform1.GetAttribute("_translate").Get()), Gf.Vec3d(1, 2, 3))
        self.assertEqual(Gf.Vec3d(*xform2.GetAttribute("_translate").Get()), Gf.Vec3d(1, 1, 1))

    # ----------------------------------------------------------------------
    async def test_gatherprototype_performance(self):
        """Exercises the GatherPrototype for the purpose of performance measurement"""
        # Test will include test_dim^3 elements
        test_dim = 5
        test_iterations = 10
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("GetPrimPaths", "omni.graph.nodes.FindPrims"),
                    ("GatherByPath", "omni.graph.nodes.GatherByPath"),
                    ("Get", "omni.graph.nodes.GetGatheredAttribute"),
                    ("Mult", "omni.graph.nodes.Multiply"),
                    ("ConstV", "omni.graph.nodes.ConstantVec3d"),
                    ("Set", "omni.graph.nodes.SetGatheredAttribute"),
                ],
                keys.CREATE_PRIMS: [
                    (f"Test_{x}_{y}_{z}", {"_translate": ("pointd[3]", [x * test_dim, y * test_dim, z * test_dim])})
                    for x in range(1, test_dim + 1)
                    for y in range(1, test_dim + 1)
                    for z in range(1, test_dim + 1)
                ],
                keys.EXPOSE_PRIMS: [
                    (og.Controller.PrimExposureType.AS_WRITABLE, f"/Test_{x}_{y}_{z}", f"Write_{x}_{y}_{z}")
                    for x in range(1, test_dim + 1)
                    for y in range(1, test_dim + 1)
                    for z in range(1, test_dim + 1)
                ],
                keys.CONNECT: [
                    ("GetPrimPaths.outputs:primPaths", "GatherByPath.inputs:primPaths"),
                    ("GatherByPath.outputs:gatherId", "Get.inputs:gatherId"),
                    ("Get.outputs:value", "Mult.inputs:a"),
                    ("ConstV.inputs:value", "Mult.inputs:b"),
                    ("Mult.outputs:product", "Set.inputs:value"),
                    ("GatherByPath.outputs:gatherId", "Set.inputs:gatherId"),
                ],
                keys.SET_VALUES: [
                    ("GetPrimPaths.inputs:namePrefix", "Test_"),
                    ("GatherByPath.inputs:attributes", "_translate"),
                    ("GatherByPath.inputs:allAttributes", False),
                    ("GatherByPath.inputs:shouldWriteBack", True),
                    ("Get.inputs:name", "_translate"),
                    ("Set.inputs:name", "_translate"),
                    ("ConstV.inputs:value", [10, 20, 30]),
                ],
            },
        )

        await controller.evaluate(graph)

        # Sanity check it did one and only one evaluation
        val = Gf.Vec3d(*controller.prim("/Test_2_1_1").GetAttribute("_translate").Get())
        expected = Gf.Vec3d(test_dim * 2 * 10, test_dim * 20, test_dim * 30)
        self.assertEqual(val, expected)
        import time

        t0 = time.time()
        for _ in range(test_iterations):
            await controller.evaluate(graph)
        dt = time.time() - t0
        print(f"{dt:3.2}s for {test_iterations} iterations of {test_dim**3} elements")

    # ----------------------------------------------------------------------
    async def test_gatherprototype_hydra_fastpath(self):
        """Tests the Gather Prototype Hydra attribute creation feature"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        cube1 = UsdGeom.Xformable(ogts.create_cube(stage, "Test_Xform1", (1, 1, 1)))
        cube1.AddTranslateOp().Set(Gf.Vec3d(1, 2, 3))
        cube1.AddRotateXYZOp().Set((90, 0, 0))
        cube2 = UsdGeom.Xformable(ogts.create_cube(stage, "Test_Xform2", (1, 1, 1)))
        cube2.AddTranslateOp().Set(Gf.Vec3d(4, 5, 6))
        cube2.AddRotateZOp().Set(90)

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("GatherByPath", "omni.graph.nodes.GatherByPath"),
                    ("GetWorldPos", "omni.graph.nodes.GetGatheredAttribute"),
                    ("GetWorldOrient", "omni.graph.nodes.GetGatheredAttribute"),
                ],
                keys.CONNECT: [
                    ("GatherByPath.outputs:gatherId", "GetWorldPos.inputs:gatherId"),
                    ("GatherByPath.outputs:gatherId", "GetWorldOrient.inputs:gatherId"),
                ],
                keys.SET_VALUES: [
                    ("GatherByPath.inputs:primPaths", ["/Test_Xform1", "/Test_Xform2"]),
                    ("GatherByPath.inputs:hydraFastPath", "World"),
                    ("GetWorldPos.inputs:name", "_worldPosition"),
                    ("GetWorldOrient.inputs:name", "_worldOrientation"),
                ],
            },
        )
        await controller.evaluate(graph)

        # Check that hydra attribs were added and that they match what we expect
        vals = og.Controller.get(controller.attribute("outputs:value", "GetWorldPos", graph))
        self.assertEqual(len(vals), 2)
        self.assertEqual(Gf.Vec3d(*vals[1]), Gf.Vec3d([4, 5, 6]))

        vals = og.Controller.get(controller.attribute("outputs:value", "GetWorldOrient", graph))
        self.assertEqual(len(vals), 2)
        val = vals[1].tolist()
        got = Gf.Quatd(val[2], val[0], val[1], val[2])
        expected = Gf.Transform(cube2.GetLocalTransformation()).GetRotation().GetQuat()
        self.assertAlmostEqual(got.GetReal(), expected.GetReal())
        for e, e2 in zip(got.GetImaginary(), expected.GetImaginary()):
            self.assertAlmostEqual(e, e2)

        # switch to Local mode
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("GetLocalMatrix", "omni.graph.nodes.GetGatheredAttribute"),
                keys.CONNECT: ("GatherByPath.outputs:gatherId", "GetLocalMatrix.inputs:gatherId"),
                keys.SET_VALUES: [
                    ("GatherByPath.inputs:hydraFastPath", "Local"),
                    ("GetLocalMatrix.inputs:name", "_localMatrix"),
                ],
            },
        )
        await controller.evaluate(graph)

        vals = og.Controller.get(controller.attribute("outputs:value", "GetLocalMatrix", graph))
        self.assertEqual(len(vals), 2)
        got = Gf.Matrix4d(*vals[1].tolist())
        expected = cube2.GetLocalTransformation()
        self.assertEqual(got, expected)

    # ----------------------------------------------------------------------
    async def test_gatherprototype_writeback(self):
        """Tests the Gather Prototype Writeback to USD feature

        Curently fails on linux only... Temporarily comment out during investigation...

        controller = og.Controller()
        keys = og.Controller.Keys

        (
            graph,
            _,
            (writeback_prim, no_writeback_prim),
            _,
        ) = controller.edit(self.TEST_GRAPH_PATH, {
            keys.CREATE_NODES: [
                ("GatherByPathWriteback", "omni.graph.nodes.GatherByPath"),
                ("GatherByPathNoWriteback", "omni.graph.nodes.GatherByPath"),
                ("SetWriteback", "omni.graph.nodes.SetGatheredAttribute"),
                ("SetNoWriteback", "omni.graph.nodes.SetGatheredAttribute")
            ],
            keys.CREATE_PRIMS: [
                ("XformWriteback", {"_translate": ("pointd[3]", [0, 0, 0])}),
                ("XformNoWriteback", {"_translate": ("pointd[3]", [0, 0, 0])}),
            ],
            keys.EXPOSE_PRIMS: (og.Controller.PrimExposureType.AS_WRITABLE, "/XformWriteback", "WritePrim"),
            keys.CONNECT: [
                ("GatherByPathWriteback.outputs:gatherId", "SetWriteback.inputs:gatherId"),
                ("GatherByPathNoWriteback.outputs:gatherId", "SetNoWriteback.inputs:gatherId")
            ],
            keys.SET_VALUES: [
                ("GatherByPathWriteback.inputs:primPaths", ["/XformWriteback"]),
                ("GatherByPathNoWriteback.inputs:primPaths", ["/XformNoWriteback"]),
                ("GatherByPathWriteback.inputs:attributes", "_translate"),
                ("GatherByPathWriteback.inputs:allAttributes", False),
                ("GatherByPathWriteback.inputs:shouldWriteBack", True),
                ("GatherByPathNoWriteback.inputs:attributes", "_translate"),
                ("GatherByPathNoWriteback.inputs:allAttributes", False),
                ("GatherByPathNoWriteback.inputs:shouldWriteBack", False),
                ("SetWriteback.inputs:name", "_translate"),
                ("SetWriteback.inputs:value", {"type": "double[3][]", "value": [[1, 2, 3]]}),
                ("SetNoWriteback.inputs:name", "_translate"),
                ("SetNoWriteback.inputs:value", {"type": "double[3][]", "value": [[1, 2, 3]]})
            ],
        })
        await controller.evaluate(graph)

        # Check that the attribs in USD were only changed for the specified prim
        write_back_attrib = writeback_prim.GetAttribute("_translate")
        no_write_back_attrib = no_writeback_prim.GetAttribute("_translate")
        self.assertEqual(write_back_attrib.Get(), Gf.Vec3d(1, 2, 3))
        self.assertEqual(no_write_back_attrib.Get(), Gf.Vec3d(0, 0, 0))
        """

    # ----------------------------------------------------------------------
    async def test_gatherprototype_repeated_paths(self):
        """Tests the Gather Prototype with repeated prim paths in the input"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # Test the Getter. Getter should get two of the same data
        (graph, (_, get_attrib_node, _), (xform_prim,), _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("GatherByPath", "omni.graph.nodes.GatherByPath"),
                    ("GetAttrib", "omni.graph.nodes.GetGatheredAttribute"),
                ],
                keys.CREATE_PRIMS: ("Xform", {"_translate": ["pointd[3]", [1, 2, 3]]}),
                keys.EXPOSE_PRIMS: (og.Controller.PrimExposureType.AS_WRITABLE, "/Xform", "WritePrim"),
                keys.CONNECT: [("GatherByPath.outputs:gatherId", "GetAttrib.inputs:gatherId")],
                keys.SET_VALUES: [
                    ("GatherByPath.inputs:primPaths", ["/Xform", "/Xform"]),
                    ("GatherByPath.inputs:attributes", "_translate"),
                    ("GatherByPath.inputs:allAttributes", False),
                    ("GetAttrib.inputs:name", "_translate"),
                ],
            },
        )
        await controller.evaluate(graph)

        vals = og.Controller.get(controller.attribute("outputs:value", get_attrib_node))
        self.assertEqual(len(vals), 2)
        self.assertEqual(Gf.Vec3d(*vals[0]), Gf.Vec3d([1, 2, 3]))
        self.assertEqual(Gf.Vec3d(*vals[1]), Gf.Vec3d([1, 2, 3]))
        self.assertEqual(Gf.Vec3d(*xform_prim.GetAttribute("_translate").Get()), Gf.Vec3d(1, 2, 3))

        # Test the Setter. Both masks should change the same prim value
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("SetAttrib", "omni.graph.nodes.SetGatheredAttribute"),
                keys.CONNECT: ("GatherByPath.outputs:gatherId", "SetAttrib.inputs:gatherId"),
                keys.SET_VALUES: [
                    ("SetAttrib.inputs:name", "_translate"),
                    ("SetAttrib.inputs:mask", [True, False]),
                    ("SetAttrib.inputs:value", [[4, 5, 6], [7, 8, 9]], "double[3][]"),
                    ("GatherByPath.inputs:shouldWriteBack", True),
                ],
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(Gf.Vec3d(*xform_prim.GetAttribute("_translate").Get()), Gf.Vec3d(4, 5, 6))

        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: ("SetAttrib.inputs:mask", [False, True])})
        await controller.evaluate(graph)
        self.assertEqual(Gf.Vec3d(*xform_prim.GetAttribute("_translate").Get()), Gf.Vec3d(7, 8, 9))
