""" Tests for omnigraph instancing """

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 TestAutoInstancing(ogts.OmniGraphTestCase):
    async def setUp(self):
        await super().setUp()
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_fast_mode(True)

    # ----------------------------------------------------------------------
    async def tick(self):
        # Auto instance merging happens globaly before any compute takes place,
        #   but the decision of asking to be merged is a per-graph thing that can happen
        #   at the beggining of each graph compute, when the pipeline is already running
        # So we need 2 frames in those tests to ensure that merging actually happens
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()

    # ----------------------------------------------------------------------
    def create_simple_graph(self, graph_path, prim_path, value: int):
        controller = og.Controller()
        keys = og.Controller.Keys

        #  +------------+       +------------+
        #  | NewBundle  |======>| InsertAttr |
        #  +------------+       |            |
        #                   ===>|val         |
        #                   |   +------------+
        #                   |
        #                   |   +---------------+
        #                   |   | WritePrimAttr |
        #       +------+    |   |               |
        #       |Const |=======>|value          |
        #       +------+        |               |
        #                       |               |
        #       +------+        |               |
        #       |Const |=======>|path           |
        #       +------+        |               |
        #                       |               |
        # +------------+        |               |
        # |ReadVariable|=======>|name           |
        # +------------+        +---------------+
        #
        # +--------------+
        # |Tutorial State|
        # +--------------+

        # The tutorial state allows to exercise internal states when using auto instancing
        # it does not need any connection to validate that it works as expected

        # create the target prim
        stage = omni.usd.get_context().get_stage()
        prim = stage.DefinePrim(prim_path)
        prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Int).Set(0)
        prim.CreateAttribute("graph_output2", Sdf.ValueTypeNames.Int).Set(0)
        prim.CreateAttribute("graph_output3", Sdf.ValueTypeNames.Int).Set(0)
        prim.CreateAttribute("graph_output4", Sdf.ValueTypeNames.Int).Set(0)
        prim.CreateAttribute("graph_output5", Sdf.ValueTypeNames.Int).Set(0)

        # create the graph
        (graph, _, _, _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadVariable", "omni.graph.core.ReadVariable"),
                    ("Path", "omni.graph.nodes.ConstantToken"),
                    ("Value", "omni.graph.nodes.ConstantInt"),
                    ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                    ("State", "omni.graph.tutorials.State"),
                    ("NewBundle", "omni.graph.nodes.BundleConstructor"),
                    ("Insert", "omni.graph.nodes.InsertAttribute"),
                ],
                keys.CONNECT: [
                    ("ReadVariable.outputs:value", "WritePrimAttr.inputs:name"),
                    ("Path.inputs:value", "WritePrimAttr.inputs:primPath"),
                    ("Value.inputs:value", "WritePrimAttr.inputs:value"),
                    ("NewBundle.outputs_bundle", "Insert.inputs:data"),
                    ("Value.inputs:value", "Insert.inputs:attrToInsert"),
                ],
                keys.CREATE_VARIABLES: [
                    ("attrib_name", og.Type(og.BaseDataType.TOKEN)),
                ],
                keys.SET_VALUES: [
                    ("ReadVariable.inputs:variableName", "attrib_name"),
                    ("WritePrimAttr.inputs:usePath", True),
                    ("Path.inputs:value", prim_path),
                    ("Value.inputs:value", value),
                    ("Insert.inputs:outputAttrName", "value"),
                ],
            },
        )

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

        return graph

    # ----------------------------------------------------------------------
    def create_instance_graph(self, graph_path, prim_root_path, inst_count, value: int):
        controller = og.Controller()
        keys = og.Controller.Keys

        #  +------------+       +------------+
        #  | NewBundle  |======>| InsertAttr |
        #  +------------+       |            |
        #                   ===>|val         |
        #                   |   +------------+
        #                   |
        #                   |   +---------------+
        #                   |   | WritePrimAttr |
        #       +------+    |   |               |
        #       |Const |=======>|value          |
        #       +------+        |               |
        #                       |               |
        #  +------------+       |               |
        #  |GraphTarget |======>|path           |
        #  +------------+       |               |
        #                       |               |
        # +------------+        |               |
        # |ReadVariable|=======>|name           |
        # +------------+        +---------------+
        #
        # +--------------+
        # |Tutorial State|
        # +--------------+

        (graph, _, _, _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadVariable", "omni.graph.core.ReadVariable"),
                    ("Target", "omni.graph.nodes.GraphTarget"),
                    ("Value", "omni.graph.nodes.ConstantInt"),
                    ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                    ("State", "omni.graph.tutorials.State"),
                    ("NewBundle", "omni.graph.nodes.BundleConstructor"),
                    ("Insert", "omni.graph.nodes.InsertAttribute"),
                ],
                keys.CONNECT: [
                    ("ReadVariable.outputs:value", "WritePrimAttr.inputs:name"),
                    ("Target.outputs:primPath", "WritePrimAttr.inputs:primPath"),
                    ("Value.inputs:value", "WritePrimAttr.inputs:value"),
                    ("NewBundle.outputs_bundle", "Insert.inputs:data"),
                    ("Value.inputs:value", "Insert.inputs:attrToInsert"),
                ],
                keys.CREATE_VARIABLES: [
                    ("attrib_name", og.Type(og.BaseDataType.TOKEN)),
                ],
                keys.SET_VALUES: [
                    ("ReadVariable.inputs:variableName", "attrib_name"),
                    ("WritePrimAttr.inputs:usePath", True),
                    ("Value.inputs:value", value),
                    ("Insert.inputs:outputAttrName", "value"),
                ],
            },
        )

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

        # create instances
        stage = omni.usd.get_context().get_stage()
        for i in range(0, inst_count):
            prim_name = f"{prim_root_path}_{i}"
            prim = stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, graph_path)
            prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Int).Set(0)
            prim.CreateAttribute("graph_output2", Sdf.ValueTypeNames.Int).Set(0)
            prim.CreateAttribute("graph_output3", Sdf.ValueTypeNames.Int).Set(0)
            prim.CreateAttribute("graph_output4", Sdf.ValueTypeNames.Int).Set(0)
            prim.CreateAttribute("graph_output5", Sdf.ValueTypeNames.Int).Set(0)

        return graph

    # ----------------------------------------------------------------------
    async def setup_stage(self, use_usd, use_og_inst):
        stage = omni.usd.get_context().get_stage()
        inst_count = 4
        int_range = range(1, 10)
        if use_usd:
            # main template
            main_template_path = "/World/Template"
            main_template_prim = f"{main_template_path}/Prim"
            main_template_graph = f"{main_template_path}/Graph"
            if use_og_inst:
                graph = self.create_instance_graph(main_template_graph, main_template_prim, inst_count, 1)
            else:
                graph = self.create_simple_graph(main_template_graph, main_template_prim, 1)

            for i in int_range:
                smart_asset_path = f"/World/SmartAsset_{i}"
                instance_prim = stage.DefinePrim(smart_asset_path)
                instance_prim.GetReferences().AddReference("", main_template_path)

            # give a chance to the graph to be created
            await self.tick()

            # setup the overs
            for i in int_range:
                smart_asset_path = f"/World/SmartAsset_{i}"
                # keep the first value inherited from template
                if i != 1:
                    attr = stage.GetPrimAtPath(f"{smart_asset_path}/Graph/Value").GetAttribute("inputs:value")
                    attr.Set(i)
                if use_og_inst is False:
                    attr = stage.GetPrimAtPath(f"{smart_asset_path}/Graph/Path").GetAttribute("inputs:value")
                    attr.Set(f"/World/SmartAsset_{i}/Prim")
        else:
            for i in int_range:
                if use_og_inst:
                    graph = self.create_instance_graph(f"/World/TestGraph_{i}", f"/World/Prim_{i}", inst_count, i)
                else:
                    graph = self.create_simple_graph(f"/World/TestGraph_{i}", f"/World/Prim_{i}", i)

        return graph

    # ----------------------------------------------------------------------
    def validate_results_and_reset(
        self,
        use_usd: bool,
        use_og_inst: bool,
        output: str = "graph_output",
        first_sa: int = 1,
        last_sa: int = 10,
        first_og_inst: int = 0,
        last_og_inst: int = 4,
        mult: int = 1,
    ):

        if use_usd:
            graph_path = "/World/SmartAsset_{asset}/Graph"
            base_path = "/World/SmartAsset_{asset}/Prim"
            base_path_og = "/World/SmartAsset_{asset}/Prim_{inst}"
        else:
            graph_path = "/World/TestGraph_{asset}"
            base_path = "/World/Prim_{asset}"
            base_path_og = "/World/Prim_{asset}_{inst}"

        # for internals states checks:
        #   all graph instances must have a different hundreds (ie. each has its own state)
        #   all graph instances must have a same mode 100 (ie. each one has executed against its own state)
        state_hundreds_last_val = -1
        state_last_mod100 = -1

        stage = omni.usd.get_context().get_stage()
        found_master_sa = False
        for i in range(first_sa, last_sa):
            graph = og.Controller.graph(graph_path.format(asset=i))
            state_out = og.Controller.attribute(graph_path.format(asset=i) + "/State.outputs:monotonic")
            insert_node = og.Controller.node(graph_path.format(asset=i) + "/Insert")
            self.assertTrue(graph.is_auto_instanced())
            inst_count = graph.get_instance_count()
            if use_og_inst:
                if found_master_sa:
                    self.assertTrue(inst_count == 4)
                else:
                    if inst_count != 4:
                        self.assertTrue(inst_count == 36)
                        found_master_sa = True
                for j in range(first_og_inst, last_og_inst):
                    prim_path = base_path_og.format(asset=i, inst=j)
                    attr = stage.GetPrimAtPath(prim_path).GetAttribute(output)
                    self.assertEqual(attr.Get(), i * mult)
                    attr.Set(0)
                    state_val = state_out.get(instance=j)
                    self.assertNotEqual(state_val // 100, state_hundreds_last_val)
                    state_hundreds_last_val = state_val // 100
                    if state_last_mod100 != -1:
                        self.assertEqual(state_val % 100, state_last_mod100)
                    state_last_mod100 = state_val % 100
                    bundle_out = graph.get_default_graph_context().get_output_bundle(insert_node, "outputs_data", j)
                    bundle_attrib = bundle_out.get_attribute_by_name("value")
                    self.assertEqual(bundle_attrib.get(), i * mult)
            else:
                if found_master_sa:
                    self.assertTrue(inst_count == 0)
                else:
                    if inst_count != 0:
                        self.assertTrue(inst_count == 9)
                        found_master_sa = True
                prim_path = base_path.format(asset=i)
                attr = stage.GetPrimAtPath(prim_path).GetAttribute(output)
                self.assertEqual(attr.Get(), i * mult)
                attr.Set(0)
                state_val = state_out.get()
                self.assertNotEqual(state_val // 100, state_hundreds_last_val)
                state_hundreds_last_val = state_val // 100
                if state_last_mod100 != -1:
                    self.assertEqual(state_val % 100, state_last_mod100)
                state_last_mod100 = state_val % 100
                bundle_out = graph.get_default_graph_context().get_output_bundle(insert_node, "outputs_data")
                bundle_attrib = bundle_out.get_attribute_by_name("value")
                self.assertEqual(bundle_attrib.get(), i * mult)

    # ----------------------------------------------------------------------
    async def test_auto_instance_no_usd_without_og_instancing(self):
        # Tests that we get a vectorized compute in smart asset situation

        await self.setup_stage(False, False)
        await self.tick()
        self.validate_results_and_reset(False, False)

    # ----------------------------------------------------------------------
    async def test_auto_instance_no_usd_with_og_instancing(self):
        # Tests that we get a vectorized compute in smart asset situation where asset have OG instancing

        await self.setup_stage(False, True)
        await self.tick()
        self.validate_results_and_reset(False, True)

    # ----------------------------------------------------------------------
    async def test_auto_instance_usd_without_og_instancing(self):
        # Tests that we get a vectorized compute in smart asset situation

        await self.setup_stage(True, False)
        await self.tick()
        self.validate_results_and_reset(True, False)

    # ----------------------------------------------------------------------
    async def test_auto_instance_usd_with_og_instancing(self):
        # Tests that we get a vectorized compute in smart asset situation where asset have OG instancing

        await self.setup_stage(True, True)
        await self.tick()
        self.validate_results_and_reset(True, True)

    # ----------------------------------------------------------------------
    async def test_auto_instance_variable_without_og_instancing(self):
        # Tests that variable changes are taken into account

        template = await self.setup_stage(True, False)

        # set the template to write to 2nd output
        var = template.find_variable("attrib_name")
        og.Controller.set_variable_default_value(var, "graph_output2")
        await self.tick()
        self.validate_results_and_reset(True, False, "graph_output2")

        # change them 1 by 1 to write to 1st output
        for i in range(1, 10):
            graph = og.Controller.graph(f"/World/SmartAsset_{i}/Graph")
            var = graph.find_variable("attrib_name")
            og.Controller.set_variable_default_value(var, "graph_output3")
            await self.tick()
            self.validate_results_and_reset(True, False, "graph_output3", 1, i + 1)
            if i != 9:
                self.validate_results_and_reset(True, False, "graph_output2", i + 1, 10)

    # ----------------------------------------------------------------------
    async def test_auto_instance_variable_with_og_instancing(self):
        # Tests that variable changes are taken into account

        stage = omni.usd.get_context().get_stage()
        template = await self.setup_stage(True, True)

        await self.tick()
        self.validate_results_and_reset(True, True, "graph_output")

        # set the template to write to out2
        template_var = template.find_variable("attrib_name")
        og.Controller.set_variable_default_value(template_var, "graph_output2")
        await self.tick()
        self.validate_results_and_reset(True, True, "graph_output2")

        # change first half of the OG template instances to write to out3
        for i in range(0, 2):
            prim_path = f"/World/Template/Prim_{i}"
            prim = stage.GetPrimAtPath(prim_path)
            variable = prim.CreateAttribute("graph:variable:attrib_name", Sdf.ValueTypeNames.Token)
            variable.Set("graph_output3")

        await self.tick()

        self.validate_results_and_reset(True, True, "graph_output3", 1, 10, 0, 2)
        self.validate_results_and_reset(True, True, "graph_output2", 1, 10, 2, 4)

        # override half SA/Graph to write to out4
        for i in range(1, 5):
            graph = og.Controller.graph(f"/World/SmartAsset_{i}/Graph")
            var = graph.find_variable("attrib_name")
            og.Controller.set_variable_default_value(var, "graph_output4")

        await self.tick()

        # COMES FROM: Template/OG_inst
        self.validate_results_and_reset(True, True, "graph_output3", 1, 5, 0, 2)
        # COMES FROM: SA/Graph
        self.validate_results_and_reset(True, True, "graph_output4", 1, 5, 2, 4)
        # COMES FROM: Template/OG_inst
        self.validate_results_and_reset(True, True, "graph_output3", 5, 10, 0, 2)
        # COMES FROM: Template/Graph
        self.validate_results_and_reset(True, True, "graph_output2", 5, 10, 2, 4)

        # Apply SA/OG_inst override to all known provenance
        for i in [1, 2, 5, 6]:
            for j in [0, 2]:
                prim_path = f"/World/SmartAsset_{i}/Prim_{j}"
                prim = stage.GetPrimAtPath(prim_path)
                variable = prim.CreateAttribute("graph:variable:attrib_name", Sdf.ValueTypeNames.Token)
                variable.Set("graph_output5")

        await self.tick()

        # COMES FROM: SA/OG_inst
        self.validate_results_and_reset(True, True, "graph_output5", 1, 3, 0, 1)
        # COMES FROM: Template/OG_inst
        self.validate_results_and_reset(True, True, "graph_output3", 1, 3, 1, 2)
        # COMES FROM: SA/OG_inst
        self.validate_results_and_reset(True, True, "graph_output5", 1, 3, 2, 3)
        # COMES FROM: SA/Graph
        self.validate_results_and_reset(True, True, "graph_output4", 1, 3, 3, 4)
        # COMES FROM: Template/OG_inst
        self.validate_results_and_reset(True, True, "graph_output3", 3, 5, 0, 2)
        # COMES FROM: SA/Graph
        self.validate_results_and_reset(True, True, "graph_output4", 3, 5, 2, 4)
        # COMES FROM: SA/OG_inst
        self.validate_results_and_reset(True, True, "graph_output5", 5, 7, 0, 1)
        # COMES FROM: Template/OG_inst
        self.validate_results_and_reset(True, True, "graph_output3", 5, 7, 1, 2)
        # COMES FROM: SA/OG_inst
        self.validate_results_and_reset(True, True, "graph_output5", 5, 7, 2, 3)
        # COMES FROM: Template/Graph
        self.validate_results_and_reset(True, True, "graph_output2", 5, 7, 3, 4)
        # COMES FROM: Template/OG_inst
        self.validate_results_and_reset(True, True, "graph_output3", 7, 10, 0, 2)
        # COMES FROM: Template/Graph
        self.validate_results_and_reset(True, True, "graph_output2", 7, 10, 2, 4)

    # ----------------------------------------------------------------------
    async def test_auto_instance_value_without_og_instancing(self):
        # Tests that value changes are taken into account
        stage = omni.usd.get_context().get_stage()
        await self.setup_stage(True, False)
        await self.tick()
        self.validate_results_and_reset(True, False)

        for i in range(1, 5):
            smart_asset_path = f"/World/SmartAsset_{i}"
            attr = stage.GetPrimAtPath(f"{smart_asset_path}/Graph/Value").GetAttribute("inputs:value")
            attr.Set(i * 10)

        await self.tick()
        self.validate_results_and_reset(True, False, "graph_output", 1, 5, mult=10)
        self.validate_results_and_reset(True, False, "graph_output", 5, 10)

        # same, but through OG directly
        for i in range(1, 5):
            smart_asset_path = f"/World/SmartAsset_{i}"
            attr = og.Controller.attribute(f"{smart_asset_path}/Graph/Value.inputs:value")
            attr.set(i * 20)

        await self.tick()
        self.validate_results_and_reset(True, False, "graph_output", 1, 5, mult=20)
        self.validate_results_and_reset(True, False, "graph_output", 5, 10)

    # ----------------------------------------------------------------------
    async def test_auto_instance_value_with_og_instancing(self):
        # Tests that a value change is taken into account
        stage = omni.usd.get_context().get_stage()
        await self.setup_stage(True, True)
        await self.tick()
        self.validate_results_and_reset(True, True)

        for i in range(1, 5):
            smart_asset_path = f"/World/SmartAsset_{i}"
            attr = stage.GetPrimAtPath(f"{smart_asset_path}/Graph/Value").GetAttribute("inputs:value")
            attr.Set(i * 10)

        await self.tick()
        self.validate_results_and_reset(True, True, "graph_output", 1, 5, mult=10)
        self.validate_results_and_reset(True, True, "graph_output", 5, 10)

        # same, but through OG directly
        for i in range(1, 5):
            smart_asset_path = f"/World/SmartAsset_{i}"
            attr = og.Controller.attribute(f"{smart_asset_path}/Graph/Value.inputs:value")
            attr.set(i * 20)

        await self.tick()
        self.validate_results_and_reset(True, True, "graph_output", 1, 5, mult=20)
        self.validate_results_and_reset(True, True, "graph_output", 5, 10)

    # ----------------------------------------------------------------------
    def _check_value_and_reset(self, val1, val2, with_inst):
        stage = omni.usd.get_context().get_stage()

        states_out = [
            og.Controller.attribute("/World/Graph1/State.outputs:monotonic"),
            og.Controller.attribute("/World/Graph2/State.outputs:monotonic"),
        ]
        # for internals states checks:
        #   all graph instances must have a different hundreds (ie. each has its own state)
        #   all graph instances must have a same mode 100 (ie. each one has executed against its own state)
        state_hundreds_last_val = -1
        state_last_mod100 = -1

        if with_inst:
            for i in range(0, 4):
                attr = stage.GetPrimAtPath(f"/World/Prim1_{i}").GetAttribute("graph_output")
                self.assertEqual(attr.Get(), val1)
                attr.Set(0)
                attr = stage.GetPrimAtPath(f"/World/Prim2_{i}").GetAttribute("graph_output")
                self.assertEqual(attr.Get(), val2)
                attr.Set(0)
                for state_out in states_out:
                    state_val = state_out.get(instance=i)
                    self.assertNotEqual(state_val // 100, state_hundreds_last_val)
                    state_hundreds_last_val = state_val // 100
                    if state_last_mod100 != -1:
                        self.assertEqual(state_val % 100, state_last_mod100)
                    state_last_mod100 = state_val % 100
                # bundle for graph 1
                graph = og.Controller.graph("/World/Graph1")
                insert_node = og.Controller.node("/World/Graph1/Insert")
                bundle_out = graph.get_default_graph_context().get_output_bundle(insert_node, "outputs_data", i)
                bundle_attrib = bundle_out.get_attribute_by_name("value")
                self.assertEqual(bundle_attrib.get(), val1)
                # bundle for graph 2
                graph = og.Controller.graph("/World/Graph2")
                insert_node = og.Controller.node("/World/Graph2/Insert")
                bundle_out = graph.get_default_graph_context().get_output_bundle(insert_node, "outputs_data", i)
                bundle_attrib = bundle_out.get_attribute_by_name("value")
                self.assertEqual(bundle_attrib.get(), val2)
        else:
            attr = stage.GetPrimAtPath("/World/Prim1").GetAttribute("graph_output")
            self.assertEqual(attr.Get(), val1)
            attr.Set(0)
            attr = stage.GetPrimAtPath("/World/Prim2").GetAttribute("graph_output")
            self.assertEqual(attr.Get(), val2)
            attr.Set(0)
            for state_out in states_out:
                state_val = state_out.get()
                self.assertNotEqual(state_val // 100, state_hundreds_last_val)
                state_hundreds_last_val = state_val // 100
                if state_last_mod100 != -1:
                    self.assertEqual(state_val % 100, state_last_mod100)
                state_last_mod100 = state_val % 100
            # bundle for graph 1
            graph = og.Controller.graph("/World/Graph1")
            insert_node = og.Controller.node("/World/Graph1/Insert")
            bundle_out = graph.get_default_graph_context().get_output_bundle(insert_node, "outputs_data")
            bundle_attrib = bundle_out.get_attribute_by_name("value")
            self.assertEqual(bundle_attrib.get(), val1)
            # bundle for graph 2
            graph = og.Controller.graph("/World/Graph2")
            insert_node = og.Controller.node("/World/Graph2/Insert")
            bundle_out = graph.get_default_graph_context().get_output_bundle(insert_node, "outputs_data")
            bundle_attrib = bundle_out.get_attribute_by_name("value")
            self.assertEqual(bundle_attrib.get(), val2)

    # ----------------------------------------------------------------------
    async def _change_check_reset(self, val1, val2, with_inst, with_usd):
        if with_usd:
            stage = omni.usd.get_context().get_stage()
            val_g1 = stage.GetPrimAtPath("/World/Graph1/Value").GetAttribute("inputs:value")
            val_g2 = stage.GetPrimAtPath("/World/Graph2/Value").GetAttribute("inputs:value")
            val_g1.Set(val1)
            val_g2.Set(val2)
        else:
            val_g1 = og.Controller.attribute("/World/Graph1/Value.inputs:value")
            val_g2 = og.Controller.attribute("/World/Graph2/Value.inputs:value")
            val_g1.set(val1)
            val_g2.set(val2)

        await self.tick()
        self._check_value_and_reset(val1, val2, with_inst)

    # ----------------------------------------------------------------------
    async def _instantiation_test(self, use_og_inst):

        # create the graphs
        if use_og_inst:
            inst_count = 4
            graph1 = self.create_instance_graph("/World/Graph1", "/World/Prim1", inst_count, 1)
            graph2 = self.create_instance_graph("/World/Graph2", "/World/Prim2", inst_count, 2)
        else:
            inst_count = 0
            graph1 = self.create_simple_graph("/World/Graph1", "/World/Prim1", 1)
            graph2 = self.create_simple_graph("/World/Graph2", "/World/Prim2", 2)

        # add a variable to the first one
        graph1.create_variable("_tmp_", og.Type(og.BaseDataType.TOKEN))

        # run a frame
        await self.tick()

        # check that the call was not vectorized
        self.assertFalse(graph1.is_auto_instanced())
        self.assertTrue(graph1.get_instance_count() == inst_count)
        self.assertFalse(graph2.is_auto_instanced())
        self.assertTrue(graph2.get_instance_count() == inst_count)

        # check results
        self._check_value_and_reset(1, 2, use_og_inst)

        # change values
        await self._change_check_reset(3, 4, use_og_inst, True)
        await self._change_check_reset(5, 6, use_og_inst, False)

        # add the same variable to graph2
        graph2.create_variable("_tmp_", og.Type(og.BaseDataType.TOKEN))

        # run a frame
        await self.tick()

        # check that the call is now vectorized
        self.assertTrue(graph1.is_auto_instanced())
        self.assertTrue(graph2.is_auto_instanced())
        ic1 = graph1.get_instance_count()
        ic2 = graph2.get_instance_count()
        if use_og_inst:
            self.assertTrue(
                (ic1 == 2 * inst_count and ic2 == inst_count) or (ic1 == inst_count and ic2 == 2 * inst_count)
            )
        else:
            self.assertTrue((ic1 == 2 and ic2 == 0) or (ic1 == 0 and ic2 == 2))

        # check results
        self._check_value_and_reset(5, 6, use_og_inst)

        # change values
        await self._change_check_reset(7, 8, use_og_inst, True)
        await self._change_check_reset(9, 10, use_og_inst, False)

        # force each graph to not use auto instancing
        for g in [graph1, graph2]:
            g.set_auto_instancing_allowed(False)
            await self.tick()

            # check that the call was not vectorized
            self.assertFalse(graph1.is_auto_instanced())
            self.assertTrue(graph1.get_instance_count() == inst_count)
            self.assertFalse(graph2.is_auto_instanced())
            self.assertTrue(graph2.get_instance_count() == inst_count)

            # check that it works
            # this also validates that the graph/bundle data has been properly migrated between master/slave
            self._check_value_and_reset(9, 10, use_og_inst)

            # change values
            await self._change_check_reset(11, 12, use_og_inst, True)
            await self._change_check_reset(13, 14, use_og_inst, False)

            # activate back auto instancing
            g.set_auto_instancing_allowed(True)
            await self.tick()

            # check that the call is now vectorized
            self.assertTrue(graph1.is_auto_instanced())
            self.assertTrue(graph2.is_auto_instanced())
            ic1 = graph1.get_instance_count()
            ic2 = graph2.get_instance_count()
            if use_og_inst:
                self.assertTrue(
                    (ic1 == 2 * inst_count and ic2 == inst_count) or (ic1 == inst_count and ic2 == 2 * inst_count)
                )
            else:
                self.assertTrue((ic1 == 2 and ic2 == 0) or (ic1 == 0 and ic2 == 2))

            # check that it works
            self._check_value_and_reset(13, 14, use_og_inst)

            # change values
            await self._change_check_reset(15, 16, use_og_inst, True)
            # back to initial value (for next loop or next test)
            await self._change_check_reset(9, 10, use_og_inst, False)

        # add a new variable to graph1
        graph1.create_variable("_tmp2_", og.Type(og.BaseDataType.TOKEN))

        # run a frame
        await self.tick()

        # check that the call is now back to non-vectorized
        self.assertFalse(graph1.is_auto_instanced())
        self.assertTrue(graph1.get_instance_count() == inst_count)
        self.assertFalse(graph2.is_auto_instanced())
        self.assertTrue(graph2.get_instance_count() == inst_count)

        # check results
        self._check_value_and_reset(9, 10, use_og_inst)

        # change values
        await self._change_check_reset(11, 12, use_og_inst, True)
        await self._change_check_reset(13, 14, use_og_inst, False)

    # ----------------------------------------------------------------------
    async def test_auto_instance_instantiation_without_og_instancing(self):
        # Tests that a graph modification can make it an auto instance
        await self._instantiation_test(False)

    # ----------------------------------------------------------------------
    async def test_auto_instance_instantiation_with_og_instancing(self):
        # Tests that a graph modification can (un)make it an auto instance
        await self._instantiation_test(True)

    # ----------------------------------------------------------------------
    async def test_auto_instance_signature_same_attrib(self):
        # Tests that 2 graphs that have a node with same attribute names don't get merged
        # the BooleanOr and BoolenAnd node has the same set of attributes:
        # inputs:a, inputs:b and outputs:result
        # 2 different graphs with the same topology but different nodes shouldn't get merged
        controller = og.Controller()
        keys = og.Controller.Keys

        # create the graphs
        (graph1, _, _, _) = controller.edit(
            "/World/TestGraph1",
            {
                keys.CREATE_NODES: [
                    ("/World/TestGraph1/Operation", "omni.graph.nodes.BooleanOr"),
                ],
                keys.SET_VALUES: [
                    ("/World/TestGraph1/Operation.inputs:a", {"type": "bool", "value": False}),
                    ("/World/TestGraph1/Operation.inputs:b", {"type": "bool", "value": True}),
                ],
            },
        )
        (graph2, _, _, _) = controller.edit(
            "/World/TestGraph2",
            {
                keys.CREATE_NODES: [
                    ("/World/TestGraph2/Operation", "omni.graph.nodes.BooleanAnd"),
                ],
                keys.SET_VALUES: [
                    ("/World/TestGraph2/Operation.inputs:a", {"type": "bool", "value": False}),
                    ("/World/TestGraph2/Operation.inputs:b", {"type": "bool", "value": True}),
                ],
            },
        )

        # run a frame
        await self.tick()

        # check results
        self.assertFalse(graph1.is_auto_instanced())
        self.assertFalse(graph2.is_auto_instanced())
        self.assertEqual(controller.attribute("/World/TestGraph1/Operation.outputs:result").get(), True)
        self.assertEqual(controller.attribute("/World/TestGraph2/Operation.outputs:result").get(), False)

    # ----------------------------------------------------------------------
    async def test_auto_instance_signature_different_connection(self):
        # Tests that 2 graphs that have the same set of attributes with different connections don't get merged

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

        # create the graphs
        (graph1, _, _, _) = controller.edit(
            "/World/TestGraph1",
            {
                keys.CREATE_NODES: [
                    ("Sub", "omni.graph.nodes.Subtract"),
                    ("C1", "omni.graph.nodes.ConstantDouble"),
                    ("C2", "omni.graph.nodes.ConstantDouble"),
                ],
                keys.CONNECT: [
                    ("C1.inputs:value", "Sub.inputs:a"),
                    ("C2.inputs:value", "Sub.inputs:b"),
                ],
                keys.SET_VALUES: [
                    ("C1.inputs:value", 10),
                    ("C2.inputs:value", 5),
                ],
            },
        )
        # connected the other way around
        (graph2, _, _, _) = controller.edit(
            "/World/TestGraph2",
            {
                keys.CREATE_NODES: [
                    ("/World/TestGraph2/Sub", "omni.graph.nodes.Subtract"),
                    ("/World/TestGraph2/C1", "omni.graph.nodes.ConstantDouble"),
                    ("/World/TestGraph2/C2", "omni.graph.nodes.ConstantDouble"),
                ],
                keys.CONNECT: [
                    ("/World/TestGraph2/C1.inputs:value", "/World/TestGraph2/Sub.inputs:b"),
                    ("/World/TestGraph2/C2.inputs:value", "/World/TestGraph2/Sub.inputs:a"),
                ],
                keys.SET_VALUES: [
                    ("/World/TestGraph2/C1.inputs:value", 10),
                    ("/World/TestGraph2/C2.inputs:value", 5),
                ],
            },
        )

        # run a frame
        await self.tick()

        # check results
        self.assertFalse(graph1.is_auto_instanced())
        self.assertFalse(graph2.is_auto_instanced())
        self.assertEqual(controller.attribute("/World/TestGraph1/Sub.outputs:difference").get(), 5)
        self.assertEqual(controller.attribute("/World/TestGraph2/Sub.outputs:difference").get(), -5)
