import math
import os

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


def example_deformer1(input_point, width, height, freq):
    """Implement simple deformation on the input point with the given wavelength, wave height, and frequency"""
    tx = freq * (input_point[0] - width) / width
    ty = 1.5 * freq * (input_point[1] - width) / width
    tz = height * (math.sin(tx) + math.cos(ty))
    return Gf.Vec3f(input_point[0], input_point[1], input_point[2] + tz)


class TestOmniGraphPythonExamples(ogts.OmniGraphTestCase):
    # ----------------------------------------------------------------------
    async def verify_example_deformer(self):
        """Compare the actual deformation against the expected values using example_deformer1"""
        stage = omni.usd.get_context().get_stage()
        input_grid = stage.GetPrimAtPath("/defaultPrim/inputGrid")
        self.assertEqual(input_grid.GetTypeName(), "Mesh")
        input_points_attr = input_grid.GetAttribute("points")
        self.assertIsNotNone(input_points_attr)

        output_grid = stage.GetPrimAtPath("/defaultPrim/outputGrid")
        self.assertEqual(output_grid.GetTypeName(), "Mesh")
        output_points_attr = output_grid.GetAttribute("points")
        self.assertIsNotNone(output_points_attr)

        await omni.kit.app.get_app().next_update_async()
        await og.Controller.evaluate()
        multiplier_attr = og.Controller.attribute("inputs:multiplier", "/defaultPrim/DeformerGraph/testDeformer")
        multiplier = og.Controller.get(multiplier_attr)
        # TODO: Should really be using hardcoded values from the file to ensure correct operation
        #       multiplier = 3.0
        width = 310.0
        height = multiplier * 10.0
        freq = 10.0

        # check that the deformer applies the correct deformation
        input_points = input_points_attr.Get()
        output_points = output_points_attr.Get()
        actual_points = []
        for input_point, output_point in zip(input_points, output_points):
            point = example_deformer1(input_point, width, height, freq)
            self.assertEqual(point[0], output_point[0])
            self.assertEqual(point[1], output_point[1])
            # verify that the z-coordinates computed by the test and the deformer match to three decimal places
            actual_points.append(point)
            self.assertAlmostEqual(point[2], output_point[2], 3)

        # need to wait for next update to propagate the attribute change from USD
        og.Controller.set(multiplier_attr, 0.0)
        await omni.kit.app.get_app().next_update_async()

        # check that the deformer with a zero multiplier leaves the grid undeformed
        input_points = input_points_attr.Get()
        output_points = output_points_attr.Get()
        for input_point, output_point in zip(input_points, output_points):
            self.assertEqual(input_point, output_point)

    # ----------------------------------------------------------------------
    async def test_example_deformer_python(self):
        """Tests the omnigraph.examples.deformerPy node"""
        (result, error) = await og.load_example_file("ExampleGridDeformerPython.usda")
        self.assertTrue(result, error)
        await omni.kit.app.get_app().next_update_async()
        await self.verify_example_deformer()

    # ----------------------------------------------------------------------
    async def run_python_node_creation_by_command(self, node_backed_by_usd=True):
        """
        This tests the creation of python nodes through the node definition mechanism
        We have a python node type that defines the inputs and outputs of a python node
        type.  Here, without pre-declaring these attributes in USD, we will create the
        node using its node type definition only.  The required attributes should
        automatically be created for the node.  We then wire it up to make sure it works.
        """
        graph = og.Controller.create_graph("/defaultPrim/TestGraph")

        new_node_path = "/defaultPrim/TestGraph/new_node"
        omni.kit.commands.execute(
            "CreateNode",
            graph=graph,
            node_path=new_node_path,
            node_type="omni.graph.examples.python.VersionedDeformerPy",
            create_usd=node_backed_by_usd,
        )
        node = graph.get_node(new_node_path)
        self.assertEqual(node.get_prim_path(), new_node_path)

    # ----------------------------------------------------------------------
    async def test_python_node_creation_by_command_usd(self):
        """Test that USD commands creating Python nodes also create OmniGraph nodes"""
        await self.run_python_node_creation_by_command(node_backed_by_usd=True)

    # ----------------------------------------------------------------------
    async def test_python_node_creation_by_command_no_usd(self):
        """Test that USD commands creating Python nodes also create OmniGraph nodes"""
        await self.run_python_node_creation_by_command(node_backed_by_usd=False)

    # ----------------------------------------------------------------------
    async def test_update_node_version(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        (input_grid, output_grid) = ogts.create_input_and_output_grid_meshes(stage)

        input_points_attr = input_grid.GetPointsAttr()
        output_points_attr = output_grid.GetPointsAttr()

        keys = og.Controller.Keys
        controller = og.Controller()
        (push_graph, _, _, _) = controller.edit(
            "/defaultPrim/pushGraph",
            {
                keys.CREATE_NODES: [
                    ("ReadPoints", "omni.graph.nodes.ReadPrimAttribute"),
                    ("WritePoints", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.SET_VALUES: [
                    ("ReadPoints.inputs:name", input_points_attr.GetName()),
                    ("ReadPoints.inputs:primPath", "/defaultPrim/inputGrid"),
                    ("ReadPoints.inputs:usePath", True),
                    ("WritePoints.inputs:name", output_points_attr.GetName()),
                    ("WritePoints.inputs:primPath", "/defaultPrim/outputGrid"),
                    ("WritePoints.inputs:usePath", True),
                ],
            },
        )

        test_deformer = stage.DefinePrim("/defaultPrim/pushGraph/testVersionedDeformer", "OmniGraphNode")
        test_deformer.GetAttribute("node:type").Set("omni.graph.examples.python.VersionedDeformerPy")
        test_deformer.GetAttribute("node:typeVersion").Set(0)
        deformer_input_points = test_deformer.CreateAttribute("inputs:points", Sdf.ValueTypeNames.Point3fArray)
        deformer_input_points.Set([(0, 0, 0)] * 1024)
        deformer_output_points = test_deformer.CreateAttribute("outputs:points", Sdf.ValueTypeNames.Point3fArray)
        deformer_output_points.Set([(0, 0, 0)] * 1024)
        multiplier_attr = test_deformer.CreateAttribute("inputs:multiplier", Sdf.ValueTypeNames.Float)
        multiplier_attr.Set(3)

        await controller.evaluate()

        controller.edit(
            push_graph,
            {
                keys.CONNECT: [
                    ("ReadPoints.outputs:value", deformer_input_points.GetPath().pathString),
                    (deformer_output_points.GetPath().pathString, "WritePoints.inputs:value"),
                ]
            },
        )

        # Wait for USD notice handler to construct the underlying compute graph
        await controller.evaluate()

        # This node has 2 attributes declared as part of its node definition.  We didn't create these attributes
        # above, but the code should have automatically filled them in
        version_deformer_node = push_graph.get_node("/defaultPrim/pushGraph/testVersionedDeformer")
        input_attr = version_deformer_node.get_attribute("test_input")
        output_attr = version_deformer_node.get_attribute("test_output")
        self.assertTrue(input_attr is not None and input_attr.is_valid())
        self.assertTrue(output_attr is not None and output_attr.is_valid())
        self.assertEqual(input_attr.get_name(), "test_input")
        self.assertEqual(input_attr.get_type_name(), "int")
        self.assertEqual(output_attr.get_name(), "test_output")
        self.assertEqual(output_attr.get_type_name(), "float")

        # We added an union type input to this node as well, so test that is here:
        union_attr = version_deformer_node.get_attribute("test_union_input")
        self.assertTrue(union_attr is not None and union_attr.is_valid())
        self.assertEqual(union_attr.get_name(), "test_union_input")
        self.assertEqual(union_attr.get_type_name(), "token")
        self.assertEqual(union_attr.get_extended_type(), og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION)
        self.assertEqual(union_attr.get_resolved_type().base_type, og.BaseDataType.UNKNOWN)
        union_attr.set_resolved_type(og.Type(og.BaseDataType.FLOAT))
        self.assertEqual(union_attr.get_resolved_type(), og.Type(og.BaseDataType.FLOAT))
        union_attr.set(32)
        self.assertEqual(union_attr.get(), 32)

        # Node version update happens as the nodes are constructed and checked against the node type
        multiplier_attr = test_deformer.GetAttribute("inputs:multiplier")
        self.assertFalse(multiplier_attr.IsValid())
        wavelength_attr = test_deformer.GetAttribute("inputs:wavelength")
        self.assertIsNotNone(wavelength_attr)
        self.assertEqual(wavelength_attr.GetTypeName(), Sdf.ValueTypeNames.Float)
        # attrib value should be initialized to 50
        og_attr = controller.node("/defaultPrim/pushGraph/testVersionedDeformer").get_attribute("inputs:wavelength")
        self.assertEqual(og_attr.get(), 50.0)

    # ----------------------------------------------------------------------
    async def test_example_tutorial1_finished(self):
        """Tests the deforming_text_tutorial_finished.usda file"""
        ext_dir = os.path.dirname(__file__) + ("/.." * 5)
        file_path = os.path.join(ext_dir, "data", "deforming_text_tutorial_finished.usda")
        (result, error) = await ogts.load_test_file(file_path)
        self.assertTrue(result, error)
        await omni.kit.app.get_app().next_update_async()
        await og.Controller.evaluate()

        # Verify the test graph is functional.  We take a look at the location of the sphere and the 'G' mesh.  When we
        # move the sphere it should also move the 'G'.  If the 'G' is not moved from it's orginal location we know
        # something is wrong.
        stage = omni.usd.get_context().get_stage()
        sphere = UsdGeom.Xformable(stage.GetPrimAtPath("/World/Sphere"))
        text_g_mesh = UsdGeom.Xformable(stage.GetPrimAtPath("/World/Text_G"))
        text_g_x0 = text_g_mesh.GetLocalTransformation()[3]
        # Move the sphere
        sphere.ClearXformOpOrder()
        sphere_translate_op = sphere.AddTranslateOp()
        sphere_translate_op.Set(Gf.Vec3d(0, 0, 0))
        # Check the 'G' position has changed
        await omni.kit.app.get_app().next_update_async()
        text_g_x1 = text_g_mesh.GetLocalTransformation()[3]
        self.assertNotEqual(text_g_x0, text_g_x1)

    # ----------------------------------------------------------------------
    async def test_singleton(self):
        """
        Basic idea of test: we define a node that's been tagged as singleton.  We then try to add another
        node of the same type.  It should fail and not let us.
        """
        (graph, (singleton_node,), _, _) = og.Controller.edit(
            "/TestGraph",
            {og.Controller.Keys.CREATE_NODES: ("testSingleton", "omni.graph.examples.python.TestSingleton")},
        )
        singleton_node = graph.get_node("/TestGraph/testSingleton")
        self.assertTrue(singleton_node.is_valid())

        with ogts.ExpectedError():
            (_, (other_node,), _, _) = og.Controller.edit(
                graph, {og.Controller.Keys.CREATE_NODES: ("testSingleton2", "omni.graph.examples.python.TestSingleton")}
            )
            self.assertFalse(other_node.is_valid())
            node = graph.get_node("/TestGraph/testSingleton")
            self.assertTrue(node.is_valid())
            node = graph.get_node("/TestGraph/testSingleton2")
            self.assertFalse(node.is_valid())

    # ----------------------------------------------------------------------
    async def test_set_compute_incomplete(self):
        """
        Test usage of OgnCountTo which leverages set_compute_incomplete API.
        """
        controller = og.Controller()
        keys = og.Controller.Keys
        (_, nodes, _, _,) = controller.edit(
            {"graph_path": "/World/TestGraph", "evaluator_name": "dirty_push"},
            {
                keys.CREATE_NODES: [
                    ("CountTo", "omni.graph.examples.python.CountTo"),
                ],
                keys.SET_VALUES: [
                    ("CountTo.inputs:increment", 1),
                ],
            },
        )

        count_out = controller.attribute("outputs:count", nodes[0])
        self.assertEqual(og.Controller.get(count_out), 0.0)

        await controller.evaluate()
        self.assertEqual(og.Controller.get(count_out), 1.0)

        await controller.evaluate()
        self.assertEqual(og.Controller.get(count_out), 2.0)

        await controller.evaluate()
        self.assertEqual(og.Controller.get(count_out), 3.0)

        await controller.evaluate()
        self.assertEqual(og.Controller.get(count_out), 3.0)
