import math
import re

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):
    tx = freq * (input_point[0] - width) / width
    ty = 1.5 * freq * (input_point[1] - width) / width
    return Gf.Vec3f(input_point[0], input_point[1], input_point[2] + height * (math.sin(tx) + math.cos(ty)))


def example_deformer2(input_point, threshold):
    return Gf.Vec3f(input_point[0], input_point[1], min(input_point[2], threshold))


class TestOmniGraphExamples(ogts.OmniGraphTestCase):
    async def verify_example_deformer(self, deformer_path: str = "/defaultPrim/testDeformer"):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()

        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)

        test_deformer = stage.GetPrimAtPath(deformer_path)
        multiplier_attr = test_deformer.GetAttribute("inputs:multiplier")
        wavelength_attr = test_deformer.GetAttribute("inputs:wavelength")

        multiplier = multiplier_attr.Get()
        width = 310.0
        if wavelength_attr:
            wavelength = wavelength_attr.Get()
            if wavelength:
                width = wavelength
        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()
        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
            self.assertAlmostEqual(point[2], output_point[2], 3)

        # need to wait for next update to propagate the attribute change from USD
        multiplier_attr_set = multiplier_attr.Set(0.0)
        self.assertTrue(multiplier_attr_set)
        await controller.evaluate()

        # 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 do_example_deformer(self, with_cuda):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        (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()
        (graph, _, _, _) = controller.edit(
            "/defaultPrim/PushGraph",
            {
                keys.CREATE_NODES: [
                    ("ReadPoints", "omni.graph.nodes.ReadPrimAttribute"),
                    ("WritePoints", "omni.graph.nodes.WritePrimAttribute"),
                    ("testDeformer", f"omni.graph.examples.cpp.{'Deformer1Gpu' if with_cuda else 'Deformer1'}"),
                ],
                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),
                    ("testDeformer.inputs:multiplier", 3.0),
                    ("testDeformer.inputs:wavelength", 310.0),
                ],
            },
        )

        controller.edit(
            "/defaultPrim/PushGraph",
            {
                keys.CONNECT: [
                    ("ReadPoints.outputs:value", "testDeformer.inputs:points"),
                    ("testDeformer.outputs:points", "WritePoints.inputs:value"),
                ]
            },
        )

        # Wait for USD notice handler to construct the underlying compute graph. Calls it twice because on the first
        # pass the node additions will be deferred, to ensure that the graph is completely set up when they are added.
        await controller.evaluate(graph)
        await self.verify_example_deformer("/defaultPrim/PushGraph/testDeformer")

    # ----------------------------------------------------------------------
    async def test_example_deformer_gpu(self):
        await self.do_example_deformer(True)

    # ----------------------------------------------------------------------
    async def test_example_deformer(self):
        await self.do_example_deformer(False)

    # ----------------------------------------------------------------------
    async def test_three_deformers_dirty_push(self):
        (result, error) = await og.load_example_file("ExampleThreeDeformersDirtyPush.usda")
        self.assertTrue(result, error)
        controller = og.Controller()
        graph_path = "/World/DirtyPushGraph"

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

        input_grid = stage.GetPrimAtPath("/World/inputGrid")
        self.assertEqual(input_grid.GetTypeName(), "Mesh")
        input_points_attr = input_grid.GetAttribute("points")
        self.assertIsNotNone(input_points_attr)

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

        deformer1 = stage.GetPrimAtPath(f"{graph_path}/deformer1")
        multiplier_attr1 = deformer1.GetAttribute("inputs:multiplier")
        wavelength_attr1 = deformer1.GetAttribute("inputs:wavelength")
        multiplier1 = multiplier_attr1.Get()
        width1 = 1.0
        if wavelength_attr1:
            wavelength1 = wavelength_attr1.Get()
            if wavelength1:
                width1 = wavelength1
        height1 = multiplier1 * 10.0

        deformer2 = stage.GetPrimAtPath(f"{graph_path}/deformer2")
        multiplier_attr2 = deformer2.GetAttribute("inputs:multiplier")
        wavelength_attr2 = deformer2.GetAttribute("inputs:wavelength")
        multiplier2 = multiplier_attr2.Get()
        width2 = 1.0
        if wavelength_attr2:
            wavelength2 = wavelength_attr2.Get()
            if wavelength2:
                width2 = wavelength2
        height2 = multiplier2 * 10.0

        deformer3 = stage.GetPrimAtPath(f"{graph_path}/deformer3")
        threshold_attr3 = deformer3.GetAttribute("inputs:threshold")
        threshold3 = threshold_attr3.Get()

        freq = 10.0

        # check that the chained deformers apply the correct deformation
        input_points = input_points_attr.Get()
        output_points = output_points_attr.Get()
        for input_point, output_point in zip(input_points, output_points):
            point = example_deformer1(input_point, width1, height1, freq)
            point = example_deformer1(point, width2, height2, freq)
            point = example_deformer2(point, threshold3)

            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 deformers match to four decimal places
            self.assertAlmostEqual(point[2], output_point[2], 4)

        # need to wait for next update to propagate the attribute change from USD
        multiplier_attr1.Set(0.0)
        multiplier_attr2.Set(0.0)
        threshold_attr3.Set(999.0)
        await controller.evaluate()
        await omni.kit.app.get_app().next_update_async()

        # check that the deformers leave the grid un-deformed
        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 verify_inverse_kinematics(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()

        ref_transform = Gf.Matrix4d()
        ref_transform.SetIdentity()

        # Verify initial configuration where the end effector is at the goal position
        goal_sphere = UsdGeom.Sphere(stage.GetPrimAtPath("/Goal"))
        goal_transform = goal_sphere.GetLocalTransformation()
        ref_transform.SetTranslate(Gf.Vec3d(0, 0, 2))
        self.assertEqual(goal_transform, ref_transform)

        joint_paths = ["/Hip", "/Knee", "/Ankle"]
        for i, joint_path in enumerate(joint_paths):
            joint_cube = UsdGeom.Cube(stage.GetPrimAtPath(joint_path))
            joint_transform = joint_cube.GetLocalTransformation()
            ref_transform.SetTranslate(Gf.Vec3d(0, 0, 6 - 2 * i))
            self.assertEqual(joint_transform, ref_transform)

        # Move the goal outside of the reachable space of the end effector. Verify that the end effector
        # doesn't move.
        goal_translate_op = goal_sphere.AddTranslateOp()
        goal_translate_op_set = goal_translate_op.Set(Gf.Vec3d(0, 0, -2))
        self.assertTrue(goal_translate_op_set)
        await controller.evaluate()

        for i, joint_path in enumerate(joint_paths):
            joint_cube = UsdGeom.Cube(stage.GetPrimAtPath(joint_path))
            joint_transform = joint_cube.GetLocalTransformation()
            ref_transform.SetTranslate(Gf.Vec3d(0, 0, 6 - 2 * i))
            self.assertEqual(joint_transform, ref_transform)

    async def test_inverse_kinematics(self):
        (result, error) = await og.load_example_file("ExampleIK.usda")
        self.assertTrue(result, error)
        await self.verify_inverse_kinematics()

    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.cpp.VersionedDeformer")
        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("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()

        # Node version update happens as the nodes are constructed and checked against the node type
        multiplier_attr = test_deformer.GetAttribute("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_multiple_deformers(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = controller.Keys

        push_graph = stage.DefinePrim("/defaultPrim/pushGraph", "OmniGraph")
        push_graph.CreateAttribute("graph:name", Sdf.ValueTypeNames.Token).Set("pushgraph")

        input_grid = ogts.create_grid_mesh(
            stage, "/defaultPrim/inputGrid", counts=(32, 32), display_color=(0.2784314, 0.64705884, 1)
        )
        input_points_attr = input_grid.GetPointsAttr()
        output_grids = [
            ogts.create_grid_mesh(
                stage, f"/defaultPrim/outputGrid{i}", counts=(32, 32), display_color=(0.784314, 0.64705884, 0.1)
            )
            for i in range(4)
        ]

        prims_to_expose = [(og.Controller.PrimExposureType.AS_ATTRIBUTES, "/defaultPrim/inputGrid", "InputGrid")] + [
            (og.Controller.PrimExposureType.AS_WRITABLE, f"/defaultPrim/outputGrid{i}", f"OutputGrid{i}")
            for i in range(4)
        ]

        (graph, _, _, _) = controller.edit(
            "/DeformerGraph",
            {
                keys.CREATE_NODES: [(f"deformer{i}", "omni.graph.examples.cpp.Deformer1") for i in range(4)],
                keys.EXPOSE_PRIMS: prims_to_expose,
                keys.SET_VALUES: [(f"deformer{i}.inputs:multiplier", i) for i in range(4)]
                + [(f"deformer{i}.inputs:wavelength", 310.0) for i in range(4)],
            },
        )

        # Wait for USD notice handler to construct the underlying compute graph and resolve the output types
        await controller.evaluate()
        await omni.kit.app.get_app().next_update_async()
        controller.edit(
            graph,
            {
                keys.CONNECT: [("InputGrid.outputs:points", f"deformer{i}.inputs:points") for i in range(4)]
                + [(f"deformer{i}.outputs:points", f"OutputGrid{i}.inputs:points") for i in range(4)]
            },
        )
        await controller.evaluate()

        width = 310.0
        freq = 10.0

        # # Check that the deformer applies the correct deformation
        input_points = input_points_attr.Get()
        for i, output_grid in enumerate(output_grids):
            height = 10.0 * i
            output_points = output_grid.GetPointsAttr().Get()
            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
                self.assertAlmostEqual(point[2], output_point[2], 3)

    # ----------------------------------------------------------------------
    async def test_create_all_builtins(self):
        """Test state after manually creating a graph with one of each builtin node type"""
        all_builtin_types = [
            node_type_name for node_type_name in og.get_registered_nodes() if node_type_name.startswith("omni.graph")
        ]
        stage = omni.usd.get_context().get_stage()

        re_bare_name = re.compile(r"([^\.]*)$")

        def builtin_name(builtin_type):
            """Return a canonical node name for the builtin node type name"""
            builtin_match = re_bare_name.search(builtin_type)
            if builtin_match:
                return f"{builtin_match.group(1)[0].upper()}{builtin_match.group(1)[1:]}".replace(".", "")
            return f"{builtin_type.upper()}".replace(".", "")

        all_builtin_names = [builtin_name(name) for name in all_builtin_types]

        og.Controller.create_graph("/BuiltinsGraph")
        for node_name, node_type in zip(all_builtin_names, all_builtin_types):
            path = omni.usd.get_stage_next_free_path(stage, "/BuiltinsGraph/" + node_name, True)
            _ = og.GraphController.create_node(path, node_type)
            new_prim = stage.GetPrimAtPath(path)
            self.assertIsNotNone(new_prim)

        # Delete all of the nodes before the compute graph is created as some of the nodes
        # rely on more detailed initialization to avoid crashing.
        await omni.kit.stage_templates.new_stage_async()

    # ----------------------------------------------------------------------
    async def test_registered_node_types(self):
        registered_nodes = og.get_registered_nodes()
        self.assertGreater(len(registered_nodes), 0)
        # Check for a known type to see we are returning a sensible list of registered types
        self.assertEqual("omni.graph.tutorials.SimpleData" in registered_nodes, True)
