""" Tests for omnigraph instancing """

import os

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from pxr import OmniGraphSchemaTools, Sdf


# ======================================================================
class TestInstancing(ogts.OmniGraphTestCase):

    _graph_path = "/World/TestGraph"

    async def setUp(self):
        await super().setUp()
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_fast_mode(True)

    # ----------------------------------------------------------------------
    async def tick(self):
        timeline = omni.timeline.get_timeline_interface()
        timeline.play()
        await og.Controller.evaluate()
        timeline.stop()

    # ----------------------------------------------------------------------
    async def test_instances_read_variables_from_instances(self):
        """Tests that variable instances are correctly read from the graph when evaluate_targets is called"""
        controller = og.Controller()
        keys = og.Controller.Keys

        #                       +---------------+
        #                       | WritePrimAttr |
        # +------------+        |               |
        # |ReadVariable|=======>|value          |
        # +------------+        |               |
        #                       |               |
        # +------------+        |               |
        # |ReadVariable|=======>|primPath       |
        # +------------+        +---------------+
        (_, _, _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadVariable", "omni.graph.core.ReadVariable"),
                    ("ReadPrimName", "omni.graph.core.ReadVariable"),
                    ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CONNECT: [
                    ("ReadVariable.outputs:value", "WritePrimAttr.inputs:value"),
                    ("ReadPrimName.outputs:value", "WritePrimAttr.inputs:primPath"),
                ],
                keys.CREATE_VARIABLES: [
                    ("float_var", og.Type(og.BaseDataType.FLOAT)),
                    ("prim_name", og.Type(og.BaseDataType.TOKEN)),
                ],
                keys.SET_VALUES: [
                    ("ReadVariable.inputs:variableName", "float_var"),
                    ("ReadPrimName.inputs:variableName", "prim_name"),
                    ("WritePrimAttr.inputs:usePath", True),
                    ("WritePrimAttr.inputs:name", "graph_output"),
                ],
            },
        )

        int_range = range(1, 100)

        stage = omni.usd.get_context().get_stage()

        for i in int_range:
            prim_name = f"/World/Prim_{i}"
            prim = stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, self._graph_path)
            prim.CreateAttribute("graph:variable:float_var", Sdf.ValueTypeNames.Float).Set(float(i))
            prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Float).Set(float(0))
            prim.CreateAttribute("graph:variable:prim_name", Sdf.ValueTypeNames.Token).Set(prim_name)

        await self.tick()

        for i in int_range:
            prim_path = f"/World/Prim_{i}"
            val = stage.GetPrimAtPath(prim_path).GetAttribute("graph_output").Get()
            self.assertEqual(val, i)

    # ----------------------------------------------------------------------
    async def test_instancing_reports_graph_target(self):
        """Tests that graphs containing a node that uses the graph target reports the
        correct value
        """
        controller = og.Controller()
        keys = og.Controller.Keys
        stage = omni.usd.get_context().get_stage()

        prims = [stage.DefinePrim(f"/prim_{i}") for i in range(1, 100)]
        for prim in prims:
            prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Token)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim.GetPath(), self._graph_path)

        #               +----------------+
        # +--------+===>|primPath        |
        # | self   |    |  writePrimAttr |
        # +--------+===>|value           |
        #               +----------------+
        (graph, _, _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                    ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CONNECT: [
                    ("GraphTarget.outputs:primPath", "WritePrimAttr.inputs:value"),
                    ("GraphTarget.outputs:primPath", "WritePrimAttr.inputs:primPath"),
                ],
                keys.SET_VALUES: [
                    ("WritePrimAttr.inputs:usePath", True),
                    ("WritePrimAttr.inputs:name", "graph_output"),
                ],
            },
        )

        graph_prim = stage.GetPrimAtPath(self._graph_path)
        graph_prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Token)

        # validate calling get_graph_target outside of execution gives the graph path
        self.assertEqual(graph.get_default_graph_context().get_graph_target(), self._graph_path)

        # evaluate the graph, validate it returns the correct prim
        graph.evaluate()

        # evaluate the prims, validate the output write their own prim path
        await self.tick()
        for prim in prims:
            self.assertEqual(prim.GetAttribute("graph_output").Get(), str(prim.GetPath()))

    # ----------------------------------------------------------------------
    async def test_instances_edits_apply(self):
        """Tests that changes to variable instances (USD) are applied to the graph immediately"""
        controller = og.Controller()
        keys = og.Controller.Keys
        int_range = range(1, 100)

        #                       +---------------+
        #                       | WritePrimAttr |
        # +------------+        |               |
        # |ReadVariable|=======>|value          |
        # +------------+        |               |
        #                       |               |
        # +------------+        |               |
        # |GraphTarget |=======>|primPath       |
        # +------------+        +---------------+
        with ogts.ExpectedError():
            (graph, _, _, _) = controller.edit(
                self._graph_path,
                {
                    keys.CREATE_NODES: [
                        ("ReadVariable", "omni.graph.core.ReadVariable"),
                        ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                        ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                    ],
                    keys.CONNECT: [
                        ("ReadVariable.outputs:value", "WritePrimAttr.inputs:value"),
                        ("GraphTarget.outputs:primPath", "WritePrimAttr.inputs:primPath"),
                    ],
                    keys.CREATE_VARIABLES: [
                        ("float_var", og.Type(og.BaseDataType.FLOAT)),
                    ],
                    keys.SET_VALUES: [
                        ("ReadVariable.inputs:variableName", "float_var"),
                        ("WritePrimAttr.inputs:usePath", True),
                        ("WritePrimAttr.inputs:name", "graph_output"),
                    ],
                },
            )

        stage = omni.usd.get_context().get_stage()

        # avoid an error because the graph gets run as well
        stage.GetPrimAtPath(self._graph_path).CreateAttribute("graph_output", Sdf.ValueTypeNames.Float).Set(float(0))

        # create the instances, avoid the override
        for i in int_range:
            prim_name = f"/World/Prim_{i}"
            prim = stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, self._graph_path)
            prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Float).Set(float(0))

        # set the variable default value
        var = graph.find_variable("float_var")
        og.Controller.set_variable_default_value(var, 10.0)

        await self.tick()

        # validate all prims read the default value
        for i in int_range:
            prim = stage.GetPrimAtPath(f"/World/Prim_{i}")
            self.assertEqual(prim.GetAttribute("graph_output").Get(), 10.0)

        # create an override attribute on each prim, set it with an unique value
        for i in int_range:
            prim = stage.GetPrimAtPath(f"/World/Prim_{i}")
            prim.CreateAttribute("graph:variable:float_var", Sdf.ValueTypeNames.Float).Set(float(i))

        await self.tick()

        # update the attribute
        for i in int_range:
            stage.GetPrimAtPath(f"/World/Prim_{i}").GetAttribute("graph:variable:float_var").Set(float(i * i))

        await self.tick()

        # validate the updated output was read in the graph
        for i in int_range:
            prim_path = f"/World/Prim_{i}"
            val = stage.GetPrimAtPath(prim_path).GetAttribute("graph_output").Get()
            self.assertEqual(val, float(i * i))

    async def test_instances_run_from_references(self):
        """OM-48547 - tests instances run when added as references"""

        num_prims = 4
        for frame_rate in [24, 25, 30]:
            with self.subTest(frame_rate=frame_rate):
                timeline = omni.timeline.get_timeline_interface()
                timeline.set_time_codes_per_second(frame_rate)

                file_path = os.path.join(os.path.dirname(__file__), "data", "TestReferenceInstance.usda")
                stage = omni.usd.get_context().get_stage()

                # Add references to a file that contains a graph and a graph instance.
                # The graph increments an attribute on the instance called 'count' each frame
                for i in range(0, num_prims):
                    prim = stage.DefinePrim(f"/World/Root{i}")
                    prim.GetReferences().AddReference(file_path)

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

                counts = []
                for i in range(0, num_prims):
                    # validate a graph was properly created
                    attr = stage.GetAttributeAtPath(f"/World/Root{i}/Instance.count")
                    self.assertTrue(attr.IsValid())
                    counts.append(attr.Get())

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

                # validate that every graph instance was executed
                for i in range(0, num_prims):
                    attr = stage.GetAttributeAtPath(f"/World/Root{i}/Instance.count")
                    self.assertGreater(attr.Get(), counts[i])
                    counts[i] = attr.Get()

                # remove a prim along with its instances. Validates no errors are thrown
                # with the next update
                stage.RemovePrim(f"/World/Root{num_prims-1}")

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

                # validate the remaining update
                for i in range(0, num_prims - 1):
                    attr = stage.GetAttributeAtPath(f"/World/Root{i}/Instance.count")
                    self.assertGreater(attr.Get(), counts[i])

                timeline.stop()
                await omni.usd.get_context().new_stage_async()

    # ----------------------------------------------------------------------
    async def test_instances_can_be_duplicated(self):
        """Test that valdiates that an OmniGraph instance can be duplicated"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # +------------+        +---------------+
        # |GraphTarget |=======>|WriteVariable  |
        # +------------+        +---------------+
        (graph, _, _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                    ("WriteVar", "omni.graph.core.WriteVariable"),
                ],
                keys.CONNECT: [
                    ("GraphTarget.outputs:primPath", "WriteVar.inputs:value"),
                ],
                keys.CREATE_VARIABLES: [
                    ("path", og.Type(og.BaseDataType.TOKEN)),
                ],
                keys.SET_VALUES: [
                    ("WriteVar.inputs:variableName", "path"),
                ],
            },
        )

        # create an instance, and duplicate it
        stage = omni.usd.get_context().get_stage()
        prim0 = stage.DefinePrim("/World/Prim0")
        OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim0.GetPath(), self._graph_path)
        omni.kit.commands.execute("CopyPrim", path_from=prim0.GetPath(), path_to="/World/Prim1", exclusive_select=False)

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

        # validate both instances ran by checking the variable was written
        variable = graph.find_variable("path")
        context = graph.get_default_graph_context()
        self.assertEquals(variable.get(context, "/World/Prim0"), "/World/Prim0")
        self.assertEquals(variable.get(context, "/World/Prim1"), "/World/Prim1")

    # ----------------------------------------------------------------------
    async def test_instances_python(self):
        """
        Tests that python nodes reads correctly their internal state when used in instantiated graphs

        The test creates 2 instances, the first one that use override, the  other not
        We check that the 2 instances behaves differently.
        Then we swith instances "override", and we check that each has its own internal state
        """
        controller = og.Controller()
        keys = og.Controller.Keys
        override_value = -1

        #                       +--------------------+
        #                       | OgnTutorialStatePy |
        # +------------+        |                    |        +---------------+
        # |Constant    |=======>|overrideValue       |        | WritePrimAttr |
        # +------------+        |                    |        |               |
        #                       |          monotonic |=======>|value          |
        # +------------+        |                    |        |               |
        # |ReadVariable|=======>|override            |        |               |
        # +------------+        +--------------------+        |               |
        #                               +------------+        |               |
        #                               |GraphTarget |=======>|primPath       |
        #                               +------------+        +---------------+
        #
        (graph, _, _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_PRIMS: [
                    ("/World/Inst1", {"graph_output": ("Int64", -2)}),
                    ("/World/Inst2", {"graph_output": ("Int64", -2)}),
                ],
                keys.CREATE_NODES: [
                    ("ReadVariable", "omni.graph.core.ReadVariable"),
                    ("OgnTutorialStatePy", "omni.graph.tutorials.StatePy"),
                    ("Const", "omni.graph.nodes.ConstantInt64"),
                    ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                    ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CONNECT: [
                    ("ReadVariable.outputs:value", "OgnTutorialStatePy.inputs:override"),
                    ("Const.inputs:value", "OgnTutorialStatePy.inputs:overrideValue"),
                    ("GraphTarget.outputs:primPath", "WritePrimAttr.inputs:primPath"),
                    ("OgnTutorialStatePy.outputs:monotonic", "WritePrimAttr.inputs:value"),
                ],
                keys.CREATE_VARIABLES: [
                    ("var_override", og.Type(og.BaseDataType.BOOL)),
                ],
                keys.SET_VALUES: [
                    ("Const.inputs:value", override_value),
                    ("ReadVariable.inputs:variableName", "var_override"),
                    ("WritePrimAttr.inputs:usePath", True),
                    ("WritePrimAttr.inputs:name", "graph_output"),
                ],
            },
        )

        # create the instances
        stage = omni.usd.get_context().get_stage()
        prim1 = stage.GetPrimAtPath("/World/Inst1")
        OmniGraphSchemaTools.applyOmniGraphAPI(stage, "/World/Inst1", self._graph_path)
        prim2 = stage.GetPrimAtPath("/World/Inst2")
        OmniGraphSchemaTools.applyOmniGraphAPI(stage, "/World/Inst2", self._graph_path)

        # set the variable default value
        var = graph.find_variable("var_override")
        og.Controller.set_variable_default_value(var, False)

        # change the value of the variable on prim1
        prim1.CreateAttribute("graph:variable:var_override", Sdf.ValueTypeNames.Bool).Set(True)

        # run a frame
        await self.tick()

        prim1_out = prim1.GetAttribute("graph_output")
        prim2_out = prim2.GetAttribute("graph_output")

        # prim1 should have the overriden value
        self.assertEqual(prim1_out.Get(), override_value)
        # prim2 should have the internal initial state value
        self.assertEqual(prim2_out.Get(), 0)

        # run a frame
        await self.tick()

        # prim1 should still have the overriden value
        self.assertEqual(prim1_out.Get(), override_value)
        # prim2 should have the internal state value incremented
        step2 = prim2_out.Get()
        self.assertNotEqual(step2, override_value)

        # turn off override on prim 1
        prim1.GetAttribute("graph:variable:var_override").Set(False)

        # run a frame
        await self.tick()

        # prim1 should now have a fresh internal state initial value
        self.assertEqual(prim1_out.Get(), 0)
        # prim2 should have the internal state value incremented
        self.assertEqual(prim2_out.Get(), 2 * step2)

        # run a frame
        await self.tick()

        # prim1 should have the internal state value incremented
        step1 = prim1_out.Get()
        self.assertNotEqual(step1, override_value)
        # the new step should be an increment of the previous one
        self.assertEqual(step1, step2 + 1)
        # prim2 should have the internal state value incremented
        self.assertEqual(prim2_out.Get(), 3 * step2)

        # change the value of the variable on prim2
        prim2.CreateAttribute("graph:variable:var_override", Sdf.ValueTypeNames.Bool).Set(True)

        # run a frame
        await self.tick()

        # prim1 should have the internal state value incremented
        self.assertEqual(prim1_out.Get(), 2 * step1)
        # prim2 should have overriden value
        self.assertEqual(prim2_out.Get(), override_value)

        # now check that deleting an instance deletes as well its internal state
        stage.RemovePrim("/World/Inst1")

        # run a frame to apply the removal
        await self.tick()

        # re-create a fresh prim
        prim1 = stage.DefinePrim("/World/Inst1")
        OmniGraphSchemaTools.applyOmniGraphAPI(stage, "/World/Inst1", self._graph_path)
        prim1.CreateAttribute("graph_output", Sdf.ValueTypeNames.Int64).Set(-2)
        prim1_out = prim1.GetAttribute("graph_output")

        # run a frame
        await self.tick()

        # prim1 should be back to the initial state
        self.assertEqual(prim1_out.Get(), 0)
        # prim2 should have overriden value
        self.assertEqual(prim2_out.Get(), override_value)

        # run a frame
        await self.tick()

        # prim1 should have the internal state value incremented
        new_step = prim1_out.Get()
        self.assertNotEqual(new_step, override_value)
        self.assertEqual(new_step, step1 + 1)
        # prim2 should have overriden value
        self.assertEqual(prim2_out.Get(), override_value)

        # run a frame
        await self.tick()

        # prim1 should have the internal state value incremented
        self.assertEqual(prim1_out.Get(), 2 * new_step)

    # -------------------------------------------------------------------------
    async def test_instance_variables_values_persist(self):
        """
        Tests that instance variables value persist even if the backing values are removed
        from fabric
        """
        controller = og.Controller()
        keys = og.Controller.Keys

        # Create a simple graph that reads a variable and passes it through the graph
        (graph, (read_var, to_str), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadVariable", "omni.graph.core.ReadVariable"),
                    ("ToString", "omni.graph.nodes.ToString"),
                ],
                keys.CONNECT: [
                    ("ReadVariable.outputs:value", "ToString.inputs:value"),
                ],
                keys.CREATE_VARIABLES: [
                    ("var_str", og.Type(og.BaseDataType.UCHAR, 1, 1, og.AttributeRole.TEXT)),
                ],
                keys.SET_VALUES: [
                    ("ReadVariable.inputs:variableName", "var_str"),
                ],
            },
        )

        # Create instances with unique values per instance
        num_prims = 3
        stage = omni.usd.get_context().get_stage()
        variable = graph.find_variable("var_str")
        output_attr = to_str.get_attribute("outputs:converted")
        context = graph.get_default_graph_context()

        for i in range(0, num_prims):
            prim_name = f"/World/Prim_{i}"
            stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, self._graph_path)
            attr = stage.GetPrimAtPath(f"/World/Prim_{i}").CreateAttribute(
                "graph:variable:var_str", Sdf.ValueTypeNames.String
            )
            attr.Set(str(i))

        await self.tick()

        read_var_count = read_var.get_compute_count()
        to_str_count = to_str.get_compute_count()

        for i in range(0, num_prims):
            expected_value = str(i)
            self.assertEqual(variable.get(context, instance_path=f"/World/Prim_{i}"), expected_value)
            self.assertEqual(output_attr.get(instance=i), expected_value)

        # if USDRT is available use that to remove the backing properties from Fabric.
        # otherwise, use USD.
        usdrt_stage = stage
        try:
            import usdrt

            usdrt_stage = usdrt.Usd.Stage.Attach(omni.usd.get_context().get_stage_id())
        except ImportError:
            pass

        # Update the values in Fabric
        for i in range(0, num_prims):
            variable.set(context, str(i + 10), instance_path=f"/World/Prim_{i}")

        await self.tick()

        # validate the nodes updated
        self.assertGreater(read_var.get_compute_count(), read_var_count)
        self.assertGreater(to_str.get_compute_count(), to_str_count)
        read_var_count = read_var.get_compute_count()
        to_str_count = to_str.get_compute_count()

        # Remove the backing properties
        for i in range(0, num_prims):
            prim = usdrt_stage.GetPrimAtPath(f"/World/Prim_{i}")
            prim.RemoveProperty("graph:variable:var_str")

        await self.tick()

        self.assertGreater(read_var.get_compute_count(), read_var_count)
        self.assertGreater(to_str.get_compute_count(), to_str_count)

        # validate the values still exist as expected
        for i in range(0, num_prims):
            expected_value = str(i + 10)
            self.assertEqual(variable.get(context, instance_path=f"/World/Prim_{i}"), expected_value)
            self.assertEqual(output_attr.get(instance=i), expected_value)

    # -------------------------------------------------------------------------
    async def test_instance_author_existing_graph(self):
        """
        Tests that an instanced graph can be authored
        """
        controller = og.Controller()
        keys = og.Controller.Keys

        # Create a simple graph that reads a variable and passes it through the graph
        # we need an action graph here to prevent vectorization and exercise the one-by-one code path
        # this can be revisited with a push graph when we'll be able to controll the instance batch size per node
        #  (and removing the artificially added action node that forces execution)
        controller.create_graph({"graph_path": self._graph_path, "evaluator_name": "execution"})
        (graph, (_, to_str, _, _), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadVariable", "omni.graph.core.ReadVariable"),
                    ("ToString", "omni.graph.nodes.ToString"),
                    ("ForceExec", "omni.graph.test.TestAllDataTypes"),
                    ("OnTick", "omni.graph.action.OnTick"),
                ],
                keys.CONNECT: [
                    ("ReadVariable.outputs:value", "ToString.inputs:value"),
                    ("OnTick.outputs:tick", "ForceExec.inputs:a_execution"),
                    ("ToString.outputs:converted", "ForceExec.inputs:a_string"),
                ],
                keys.CREATE_VARIABLES: [
                    ("var_str", og.Type(og.BaseDataType.UCHAR, 1, 1, og.AttributeRole.TEXT)),
                ],
                keys.SET_VALUES: [
                    ("ReadVariable.inputs:variableName", "var_str"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )

        # Create instances with unique values per instance
        num_prims = 3
        stage = omni.usd.get_context().get_stage()
        variable = graph.find_variable("var_str")
        output_attr = to_str.get_attribute("outputs:converted")
        context = graph.get_default_graph_context()

        for i in range(0, num_prims):
            prim_name = f"/World/Prim_{i}"
            stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, self._graph_path)
            attr = stage.GetPrimAtPath(f"/World/Prim_{i}").CreateAttribute(
                "graph:variable:var_str", Sdf.ValueTypeNames.String
            )
            attr.Set(str(i))

        await self.tick()

        for i in range(0, num_prims):
            expected_value = str(i)
            self.assertEqual(variable.get(context, instance_path=f"/World/Prim_{i}"), expected_value)
            self.assertEqual(output_attr.get(instance=i), expected_value)

        (_, (append), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [("Append", "omni.graph.nodes.AppendString")],
                keys.CONNECT: [
                    ("ToString.outputs:converted", "Append.inputs:value"),
                    ("Append.outputs:value", "ForceExec.inputs:a_string"),
                ],
                keys.SET_VALUES: [
                    ("Append.inputs:suffix", {"value": "_suffix", "type": "string"}),
                ],
            },
        )

        await self.tick()

        output_attr = append[0].get_attribute("outputs:value")

        for i in range(0, num_prims):
            expected_value = str(i) + "_suffix"
            self.assertEqual(output_attr.get(instance=i), expected_value)
