"""Tests for the various OmniGraph commands"""

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.test
import omni.usd


class TestOmniGraphCommands(ogts.OmniGraphTestCase):

    # ----------------------------------------------------------------------
    async def test_disable_node_command(self):
        # The basic idea of the test is as follows:
        # We load up the TestTranslatingCube.usda file, which has a cube
        # that continually moves in X.  We make sure first the
        # node is working, and then we disable the node.  Once we do
        # that, we make sure cube no longer move as the node should be
        # disabled.
        (result, error) = await ogts.load_test_file("TestTranslatingCube.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph = og.get_graph_by_path("/World/moveX")

        # make sure the box and the attribute we expect are there.
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        cube_prim = stage.GetPrimAtPath("/World/Cube")
        translate_attr = cube_prim.GetAttribute("xformOp:translate")
        self.assertTrue(translate_attr.IsValid())

        # !!make a copy as the above is a reference of the data, that will change the next frame!!
        translate = translate_attr.Get()
        t0 = translate[0]

        # wait 1 frame
        await omni.kit.app.get_app().next_update_async()

        # get the updated cube position and make sure it moved
        translate = translate_attr.Get()
        t1 = translate[0]

        diff = abs(t1 - t0)
        self.assertGreater(diff, 0.1)

        # disable the node
        subtract_node = graph.get_node("/World/moveX/subtract")

        # we're piggy backing off this test to also test that the inputs and outputs of this node are
        # properly categorized (so this part is separate from the main test)
        input_attr = subtract_node.get_attribute("inputs:a")
        self.assertEqual(input_attr.get_port_type(), og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT)
        output_attr = subtract_node.get_attribute("outputs:out")
        self.assertEqual(output_attr.get_port_type(), og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT)

        og.cmds.DisableNode(node=subtract_node)

        # wait another frame, but this time make sure there is no moving of the cube
        await omni.kit.app.get_app().next_update_async()

        translate = translate_attr.Get()
        diff = abs(translate[0] - t1)
        self.assertEqual(diff, 0, 0.01)

    # ----------------------------------------------------------------------
    async def verify_example_graph_disabled(self):
        """
        basic idea behind this test: load up the translating cubes sample.  The file is set
        initially such that the graph is disabled. We'll run the command to enable
        and disable the graph a bunch of times, each time checking the position of
        the moving cube.
        """
        graph = og.get_all_graphs()[0]
        self.assertTrue(graph.is_disabled())

        translate_attr = og.Controller.usd_attribute("/World/Cube.xformOp:translate")

        # initial position of cube should have x = 0.0.  Since the graph is initially
        # disabled, this should remain the case:
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], 0.0, places=1)

        og.cmds.EnableGraph(graph=graph)
        self.assertFalse(graph.is_disabled())
        await og.Controller.evaluate()
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], -1.0, places=1)

        omni.kit.undo.undo()
        self.assertTrue(graph.is_disabled())
        await og.Controller.evaluate()
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], -1.0, places=1)

        omni.kit.undo.redo()
        self.assertFalse(graph.is_disabled())
        await og.Controller.evaluate()
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], -2.0, places=1)

        og.cmds.DisableGraph(graph=graph)
        self.assertTrue(graph.is_disabled())
        await og.Controller.evaluate()
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], -2.0, places=1)

        omni.kit.undo.undo()
        self.assertFalse(graph.is_disabled())
        await og.Controller.evaluate()
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], -3.0, places=1)

        omni.kit.undo.redo()
        self.assertTrue(graph.is_disabled())
        await og.Controller.evaluate()
        translate = translate_attr.Get()
        self.assertAlmostEqual(translate[0], -3.0, places=1)

    # ----------------------------------------------------------------------
    async def verify_example_graph_rename(self):
        """
        This test checks if renaming a graph will update the scene properly and
        keep the USD scene, the graph, and the fabric are in sync throughout
        two different renaming commands and their undos.
        """
        graph_path = "/World/moveX"
        renamed_path = f"{graph_path}Renamed"
        graph = og.Controller.graph(graph_path)

        self.assertTrue(graph.is_valid())

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

        cube_prim = stage.GetPrimAtPath("/World/Cube")
        translate_attr = cube_prim.GetAttribute("xformOp:translate")

        self.assertTrue(translate_attr.IsValid())

        # verify graph updates at the beginning
        await og.Controller.evaluate()
        translate = og.Controller.get(og.Controller.attribute(f"{graph_path}/compose.outputs:double3"))
        self.assertAlmostEqual(translate[0], -3.0, places=1)

        # Rename using MovePrimCommand
        omni.kit.commands.execute("MovePrim", path_from=graph_path, path_to=renamed_path)
        await omni.kit.app.get_app().next_update_async()
        graph = og.Controller.graph(graph_path)
        self.assertTrue(graph is None or not graph.is_valid())
        graph = og.Controller.graph(renamed_path)
        self.assertTrue(graph.is_valid())
        # verify graph update after rename
        translate = og.Controller.get(og.Controller.attribute(f"{renamed_path}/compose.outputs:double3"))
        self.assertAlmostEqual(translate[0], -4.0, places=1)

        # verify graph update after undo
        omni.kit.undo.undo()
        await omni.kit.app.get_app().next_update_async()
        graph = og.Controller.graph(renamed_path)
        self.assertTrue(graph is None or not graph.is_valid())
        graph = og.Controller.graph(graph_path)
        self.assertTrue(graph.is_valid())
        translate = og.Controller.get(og.Controller.attribute(f"{graph_path}/compose.outputs:double3"))
        self.assertAlmostEqual(translate[0], -5.0, places=1)

    # ----------------------------------------------------------------------
    async def test_rename_graph_command(self):
        (result, error) = await ogts.load_test_file("TestTranslatingCube.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        await self.verify_example_graph_rename()

    # ----------------------------------------------------------------------
    async def test_set_attr_command_v2(self):
        (result, error) = await ogts.load_test_file("TestTranslatingCube.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph = og.get_graph_by_path("/World/moveX")

        subtract_node = graph.get_node("/World/moveX/subtract")
        b_attr = subtract_node.get_attribute("inputs:b")

        og.cmds.SetAttr(attr=b_attr, value=8.8)
        value = og.Controller.get(b_attr)
        self.assertAlmostEqual(value, 8.8, places=2)

        # Test is flaky.  Disable this part for now to see if it's to blame
        # omni.kit.undo.undo()
        # value = og.Controller.get(multiplier_attr)
        # self.assertAlmostEqual(value, old_value, places=2)

        # omni.kit.undo.redo()
        # value = og.Controller.get(multiplier_attr)
        # self.assertAlmostEqual(value, 8.8, places=2)
