import omni.graph.core as og
import omni.graph.core.tests as ogts


class TestSimCycleDependencies(ogts.OmniGraphTestCase):
    # ----------------------------------------------------------------------
    async def test_compute_counts(self):

        # Create a dirty_push graph with two ConstantInt nodes feeding into
        # an Add node.
        (graph, (input_a, input_b, output, pynode), _, _) = og.Controller.edit(
            {"graph_path": "/TestGraph", "evaluator_name": "dirty_push"},
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("inputA", "omni.graph.nodes.ConstantInt"),
                    ("inputB", "omni.graph.nodes.ConstantInt"),
                    ("output", "omni.graph.nodes.Add"),
                    ("pyNode", "omni.graph.test.ComputeErrorPy"),
                ],
                og.Controller.Keys.CONNECT: [
                    ("inputA.inputs:value", "output.inputs:a"),
                    ("inputB.inputs:value", "output.inputs:b"),
                ],
                og.Controller.Keys.SET_VALUES: [("inputA.inputs:value", 3), ("inputB.inputs:value", 5)],
            },
        )

        self.assertTrue(graph is not None)
        self.assertTrue(graph.is_valid())

        # There have been no update ticks so no computation should have happened.
        self.assertEqual(input_a.get_compute_count(), 0)
        self.assertEqual(input_b.get_compute_count(), 0)
        self.assertEqual(output.get_compute_count(), 0)
        self.assertEqual(pynode.get_compute_count(), 0)

        # Wait for an update tick, then check counts.
        await og.Controller.evaluate(graph)
        self.assertEqual(input_a.get_compute_count(), 1)
        self.assertEqual(input_b.get_compute_count(), 1)
        self.assertEqual(output.get_compute_count(), 1)
        self.assertEqual(pynode.get_compute_count(), 1)

        # With nothing new dirtied, another tick should not change counts.
        await og.Controller.evaluate(graph)
        self.assertEqual(input_a.get_compute_count(), 1)
        self.assertEqual(input_b.get_compute_count(), 1)
        self.assertEqual(output.get_compute_count(), 1)
        self.assertEqual(pynode.get_compute_count(), 1)

        # This is a push graph. Pulling the output should not invoke any
        # more evaluation.
        sum_attr = output.get_attribute("outputs:sum")
        self.assertEqual(sum_attr.get(), 8)
        await og.Controller.evaluate(graph)
        self.assertEqual(input_a.get_compute_count(), 1)
        self.assertEqual(input_b.get_compute_count(), 1)
        self.assertEqual(output.get_compute_count(), 1)
        self.assertEqual(pynode.get_compute_count(), 1)

        # Change input_a. Only the counts on input_a and output should change.
        attr_a = input_a.get_attribute("inputs:value")
        attr_a.set(9)
        await og.Controller.evaluate(graph)
        self.assertEqual(sum_attr.get(), 14)
        self.assertEqual(input_a.get_compute_count(), 2)
        self.assertEqual(input_b.get_compute_count(), 1)
        self.assertEqual(output.get_compute_count(), 2)
        self.assertEqual(pynode.get_compute_count(), 1)

        # Manually increment the count on input_a.
        self.assertEqual(input_a.increment_compute_count(), 3)
        self.assertEqual(input_a.get_compute_count(), 3)
        self.assertEqual(input_b.get_compute_count(), 1)
        self.assertEqual(output.get_compute_count(), 2)
        self.assertEqual(pynode.get_compute_count(), 1)

        # Change the python node's input to make sure that its
        # compute count increments during evaluate as well.
        dummy_in_attr = pynode.get_attribute("inputs:dummyIn")
        dummy_in_attr.set(-1)
        self.assertEqual(pynode.get_compute_count(), 1)
        await og.Controller.evaluate(graph)
        self.assertEqual(pynode.get_compute_count(), 2)

        # Manually increment the python node's counter.
        self.assertEqual(pynode.increment_compute_count(), 3)
