"""Misc collection of tests for nodes in this extension"""

from typing import List

import numpy as np
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.stage_templates
import omni.kit.test
import omni.timeline
import omni.usd
from pxr import Gf, Sdf, Usd, UsdGeom


# Result of this does not match the matrix by generating attributes individually
def _generate_random_transform_matrix():
    translate = Gf.Vec3d((np.random.rand(3) * 300).tolist())
    euler = Gf.Vec3d((np.random.rand(3) * 360).tolist())
    scale = Gf.Vec3d((np.random.rand(3) * 5).tolist())
    rotation = (
        Gf.Rotation(Gf.Vec3d.ZAxis(), euler[2])
        * Gf.Rotation(Gf.Vec3d.YAxis(), euler[1])
        * Gf.Rotation(Gf.Vec3d.XAxis(), euler[0])
    )

    return Gf.Matrix4d().SetScale(scale) * Gf.Matrix4d().SetRotate(rotation) * Gf.Matrix4d().SetTranslate(translate)


def _generate_random_prim_transform_attributes(prim, use_orient, want_rotation):
    translate = Gf.Vec3d((np.random.random_integers(-100, 100, 3)).tolist())
    euler = Gf.Vec3d((np.random.random_integers(0, 720, 3)).tolist()) if want_rotation else Gf.Vec3d()
    scale = Gf.Vec3d((np.random.random_integers(1, 10, 3)).tolist())
    prim.CreateAttribute("xformOp:translate", Sdf.ValueTypeNames.Double3, False).Set(translate)
    prim.CreateAttribute("xformOp:scale", Sdf.ValueTypeNames.Double3, False).Set(scale)
    if use_orient:
        rotation = (
            (
                Gf.Rotation(Gf.Vec3d.XAxis(), euler[0])
                * Gf.Rotation(Gf.Vec3d.YAxis(), euler[1])
                * Gf.Rotation(Gf.Vec3d.ZAxis(), euler[2])
            )
            if want_rotation
            else Gf.Rotation().SetIdentity()
        )
        quat = rotation.GetQuat()
        prim.CreateAttribute("xformOp:orient", Sdf.ValueTypeNames.Quatd, False).Set(quat)
        prim.CreateAttribute("xformOpOrder", Sdf.ValueTypeNames.String, False).Set(
            ["xformOp:translate", "xformOp:orient", "xformOp:scale"]
        )
    else:
        prim.CreateAttribute("xformOp:rotateXYZ", Sdf.ValueTypeNames.Double3, False).Set(euler)
        prim.CreateAttribute("xformOpOrder", Sdf.ValueTypeNames.String, False).Set(
            ["xformOp:translate", "xformOp:rotateXYZ", "xformOp:scale"]
        )

    return translate, euler, scale


# TODO: fix the bug in OM-46746 related to non-uniform scale in parent
def _generate_test_prims(stage, suffix: str, use_orient=True, want_rotation=True) -> List[Usd.Prim]:
    # No Hierarchy case
    source1 = stage.DefinePrim(f"/World/Cube{suffix}1", "Cube")
    _generate_random_prim_transform_attributes(source1, use_orient, want_rotation)
    # source1.GetAttribute("xformOp:scale").Set(Gf.Vec3d(1, 1, 1))

    xform1 = stage.DefinePrim(f"/World/xform{suffix}1", "Cone")
    _generate_random_prim_transform_attributes(xform1, use_orient, want_rotation)
    xform1.GetAttribute("xformOp:scale").Set(Gf.Vec3d(1, 1, 1))
    source2 = stage.DefinePrim(f"{xform1.GetPath()}/Cube{suffix}2", "Cube")
    _generate_random_prim_transform_attributes(source2, use_orient, want_rotation)

    xform2 = stage.DefinePrim(f"{xform1.GetPath()}/xform{suffix}2", "Xform")
    _generate_random_prim_transform_attributes(xform2, use_orient, want_rotation)
    xform2.GetAttribute("xformOp:scale").Set(Gf.Vec3d(1, 1, 1))
    source3 = stage.DefinePrim(f"{xform2.GetPath()}/Cube{suffix}3", "Cube")
    _generate_random_prim_transform_attributes(source3, use_orient, want_rotation)

    # Negative Scale case
    source4 = stage.DefinePrim(f"/World/Cube{suffix}4", "Cube")
    _generate_random_prim_transform_attributes(source4, use_orient, want_rotation)
    source4.GetAttribute("xformOp:scale").Set(Gf.Vec3d(np.random.uniform(-2, 2, 3).tolist()))
    return [source1, source2, source3]


def _add_inputs_prim(stage, node: og.Node, src_prim: Usd.Prim, attrib: str):
    """Connect the given prim to the given Node's input attrib"""
    omni.kit.commands.execute(
        "AddRelationshipTarget",
        relationship=stage.GetPropertyAtPath(f"{node.get_prim_path()}.inputs:{attrib}"),
        target=src_prim.GetPath(),
    )


def _remove_inputs_prim(stage, node: og.Node, target_prim: Usd.Prim, attrib: str):
    """Remove the given prim to the given Node's input attrib"""
    omni.kit.commands.execute(
        "RemoveRelationshipTarget",
        relationship=stage.GetPropertyAtPath(f"{node.get_prim_path()}.inputs:{attrib}"),
        target=target_prim.GetPath(),
    )


# ======================================================================
class TestManeuverNodes(ogts.OmniGraphTestCase):
    """Run a simple unit test that exercises nodes"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    max_tries = 10  # max times to check if we have converged
    updates_per_try = 5  # number of frames to wait between checks

    async def setUp(self):
        """Set up  test environment, to be torn down when done"""
        await super().setUp()
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_start_time(1.0)
        timeline.set_end_time(10.0)
        fps = 24.0
        timeline.set_time_codes_per_second(fps)
        timeline.set_fast_mode(True)
        await omni.kit.app.get_app().next_update_async()
        np.random.seed(12345)

    # ------------------------------------------------------------------------
    async def next_try_async(self):
        for _ in range(self.updates_per_try):
            await omni.kit.app.get_app().next_update_async()

    # ------------------------------------------------------------------------
    async def test_move_to_target(self):
        """Test OgnMoveToTarget node"""
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()

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

        for use_orient_src, use_orient_target in [(True, True), (True, False)]:
            controller = og.Controller()
            sources = _generate_test_prims(stage, "Source", use_orient=use_orient_src)
            # We don't support interpolation rotation without orient
            want_rotation = use_orient_src
            targets = _generate_test_prims(stage, "Target", use_orient=use_orient_target, want_rotation=want_rotation)

            (graph, (_, move_to_target_node), _, _,) = controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CREATE_NODES: [
                        ("OnTick", "omni.graph.action.OnTick"),
                        ("MoveToTarget", "omni.graph.nodes.MoveToTarget"),
                    ],
                    keys.CONNECT: [
                        ("OnTick.outputs:tick", "MoveToTarget.inputs:execIn"),
                    ],
                    keys.SET_VALUES: [
                        ("OnTick.inputs:onlyPlayback", True),
                        ("MoveToTarget.inputs:speed", 5),
                    ],
                },
            )

            def check_progress(src: Usd.Prim, target: Usd.Prim) -> bool:
                xform_cache = UsdGeom.XformCache()
                final_transform = xform_cache.GetLocalToWorldTransform(src)
                target_transform = xform_cache.GetLocalToWorldTransform(target)
                ok = Gf.IsClose(target_transform.Transform(Gf.Vec3d()), final_transform.Transform(Gf.Vec3d()), 0.0001)

                if ok:
                    rot = target_transform.ExtractRotation().Decompose(
                        Gf.Vec3d.XAxis(), Gf.Vec3d.YAxis(), Gf.Vec3d.ZAxis()
                    )
                    rot2 = final_transform.ExtractRotation().Decompose(
                        Gf.Vec3d.XAxis(), Gf.Vec3d.YAxis(), Gf.Vec3d.ZAxis()
                    )
                    ok = Gf.IsClose(rot, rot2, 0.0001)

                    if ok:
                        ok = (
                            Gf.IsClose(
                                target_transform.Transform(Gf.Vec3d.XAxis()).GetLength(),
                                final_transform.Transform(Gf.Vec3d.XAxis()).GetLength(),
                                0.0001,
                            )
                            and Gf.IsClose(
                                target_transform.Transform(Gf.Vec3d.YAxis()).GetLength(),
                                final_transform.Transform(Gf.Vec3d.YAxis()).GetLength(),
                                0.0001,
                            )
                            and Gf.IsClose(
                                target_transform.Transform(Gf.Vec3d.ZAxis()).GetLength(),
                                final_transform.Transform(Gf.Vec3d.ZAxis()).GetLength(),
                                0.0001,
                            )
                        )
                        if ok:
                            return True
                return False

            for src in sources:
                for target in targets:
                    _add_inputs_prim(stage, move_to_target_node, src, "sourcePrim")
                    _add_inputs_prim(stage, move_to_target_node, target, "targetPrim")

                    timeline.play()
                    done = False
                    for _ in range(self.max_tries):
                        await self.next_try_async()

                        if check_progress(src, target):
                            done = True
                            break

                    self.assertTrue(done, "source never reached target")
                    _remove_inputs_prim(stage, move_to_target_node, src, "sourcePrim")
                    _remove_inputs_prim(stage, move_to_target_node, target, "targetPrim")
            omni.kit.commands.execute("DeletePrims", paths=[graph.get_path_to_graph()])

    # ------------------------------------------------------------------------
    async def test_move_to_transform(self):
        """Test OgnMoveToTransform node"""
        controller = og.Controller()
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()

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

        sources = _generate_test_prims(stage, "Source")

        target_transform = _generate_random_transform_matrix()
        (_, (_, move_to_transform_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("MoveToTransform", "omni.graph.nodes.MoveToTransform"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "MoveToTransform.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", True),
                    ("MoveToTransform.inputs:speed", 5),
                    ("MoveToTransform.inputs:target", target_transform),
                ],
            },
        )

        for src in sources:
            _add_inputs_prim(stage, move_to_transform_node, src, "prim")
            timeline.play()
            done = False
            for _ in range(self.max_tries):
                await self.next_try_async()

                src_xformable = UsdGeom.Xformable(src)
                # Move To Transform only moves to Local Transform
                final_transform = src_xformable.GetLocalTransformation()
                if Gf.IsClose(target_transform, final_transform, 0.00001):
                    done = True
                    break
            self.assertTrue(done, "source never reached target")
            _remove_inputs_prim(stage, move_to_transform_node, src, "prim")

    # # ------------------------------------------------------------------------
    async def test_scale_to_size(self):
        """Test OgnScaleToSize node"""
        controller = og.Controller()
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()

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

        sources = _generate_test_prims(stage, "Source")
        target_scale = Gf.Vec3f(np.random.rand(3).tolist())
        (_, (_, scale_to_size_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("ScaleToSize", "omni.graph.nodes.ScaleToSize"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "ScaleToSize.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", True),
                    ("ScaleToSize.inputs:speed", 5),
                    ("ScaleToSize.inputs:target", target_scale),
                ],
            },
        )

        for src in sources:
            _add_inputs_prim(stage, scale_to_size_node, src, "prim")
            timeline.play()
            done = False
            for _ in range(self.max_tries):
                await self.next_try_async()

                final_scale = UsdGeom.XformCommonAPI(src).GetXformVectors(Usd.TimeCode.Default())[2]
                if Gf.IsClose(target_scale, final_scale, 0.00001):
                    done = True
                    break
            self.assertTrue(done, "source never reached target")
            _remove_inputs_prim(stage, scale_to_size_node, src, "prim")

    # # ------------------------------------------------------------------------
    async def test_translate_to_location(self):
        """Test OgnTranslateToLocation node"""
        controller = og.Controller()
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()
        await omni.kit.app.get_app().next_update_async()

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

        sources = _generate_test_prims(stage, "Source")
        target_translation = Gf.Vec3d(np.random.rand(3).tolist())
        (_, (_, translate_to_location_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("TranslateToLocation", "omni.graph.nodes.TranslateToLocation"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "TranslateToLocation.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", True),
                    ("TranslateToLocation.inputs:speed", 5),
                    ("TranslateToLocation.inputs:target", target_translation),
                ],
            },
        )

        for src in sources:
            _add_inputs_prim(stage, translate_to_location_node, src, "prim")
            timeline.play()
            done = False
            for _ in range(self.max_tries):
                await self.next_try_async()

                final_translation = UsdGeom.XformCommonAPI(src).GetXformVectors(Usd.TimeCode.Default())[0]
                if Gf.IsClose(target_translation, final_translation, 0.00001):
                    done = True
                    break

            self.assertTrue(done, "source never reached target")
            _remove_inputs_prim(stage, translate_to_location_node, src, "prim")

    # ------------------------------------------------------------------------
    async def test_translate_to_target(self):
        """Test OgnTranslateToTarget node"""
        controller = og.Controller()
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()

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

        sources = _generate_test_prims(stage, "Source")
        targets = _generate_test_prims(stage, "Target")
        (_, (_, translate_to_target_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("TranslateToTarget", "omni.graph.nodes.TranslateToTarget"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "TranslateToTarget.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", True),
                    ("TranslateToTarget.inputs:speed", 5),
                ],
            },
        )

        for src in sources:
            for target in targets:
                _add_inputs_prim(stage, translate_to_target_node, src, "sourcePrim")
                _add_inputs_prim(stage, translate_to_target_node, target, "targetPrim")

                timeline.play()
                done = False
                for _ in range(self.max_tries):
                    await self.next_try_async()

                    xform_cache = UsdGeom.XformCache()
                    final_translation = Gf.Transform(xform_cache.GetLocalToWorldTransform(src)).GetTranslation()
                    target_translation = Gf.Transform(xform_cache.GetLocalToWorldTransform(target)).GetTranslation()
                    if Gf.IsClose(target_translation, final_translation, 0.00001):
                        done = True
                        break

                self.assertTrue(done, "source never reached target")

                _remove_inputs_prim(stage, translate_to_target_node, src, "sourcePrim")
                _remove_inputs_prim(stage, translate_to_target_node, target, "targetPrim")

    # # ------------------------------------------------------------------------
    async def test_rotate_to_orientation(self):
        """Test OgnRotateToOrientation node using orientation"""
        controller = og.Controller()
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()

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

        sources = _generate_test_prims(stage, "Source")
        target_rotation = Gf.Vec3f((np.random.rand(3) * 360).tolist())
        (_, (_, rotate_to_orientation_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("RotateToOrientation", "omni.graph.nodes.RotateToOrientation"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "RotateToOrientation.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", True),
                    ("RotateToOrientation.inputs:speed", 3),
                    ("RotateToOrientation.inputs:target", target_rotation),
                ],
            },
        )

        for src in sources:
            _add_inputs_prim(stage, rotate_to_orientation_node, src, "prim")
            timeline.play()
            done = False
            for _ in range(self.max_tries):
                await self.next_try_async()

                final_rotation = UsdGeom.XformCommonAPI(src).GetXformVectors(Usd.TimeCode.Default())[1]
                for value in zip(target_rotation, final_rotation):
                    target_angle = np.mod(value[0], 360)
                    final_angle = np.mod(value[1], 360)
                    if Gf.IsClose(target_angle, final_angle, 0.0001):
                        done = True
                if done:
                    break
            self.assertTrue(done, "source never reached target")
            _remove_inputs_prim(stage, rotate_to_orientation_node, src, "prim")

    # ------------------------------------------------------------------------
    async def test_rotate_to_target(self):
        """Test OgnRotateToTarget node using orientation"""
        controller = og.Controller()
        keys = og.Controller.Keys
        timeline = omni.timeline.get_timeline_interface()

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

        sources = _generate_test_prims(stage, "Source")
        targets = _generate_test_prims(stage, "Target")
        (_, (_, rotate_to_target_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("RotateToTarget", "omni.graph.nodes.RotateToTarget"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "RotateToTarget.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", True),
                    ("RotateToTarget.inputs:speed", 3),
                ],
            },
        )

        for src in sources:
            for target in targets:
                _add_inputs_prim(stage, rotate_to_target_node, src, "sourcePrim")
                _add_inputs_prim(stage, rotate_to_target_node, target, "targetPrim")

                done = False
                timeline.play()
                for _ in range(self.max_tries):
                    await self.next_try_async()

                    xform_cache = UsdGeom.XformCache()
                    target_translation = (
                        Gf.Transform(xform_cache.GetLocalToWorldTransform(target)).GetRotation().GetQuaternion()
                    )
                    final_translation = (
                        Gf.Transform(xform_cache.GetLocalToWorldTransform(src)).GetRotation().GetQuaternion()
                    )

                    if Gf.IsClose(target_translation.GetReal(), final_translation.GetReal(), 0.00001) and Gf.IsClose(
                        target_translation.GetImaginary(), final_translation.GetImaginary(), 0.00001
                    ):
                        done = True
                        break

                _remove_inputs_prim(stage, rotate_to_target_node, src, "sourcePrim")
                _remove_inputs_prim(stage, rotate_to_target_node, target, "targetPrim")

                self.assertTrue(done, "source never reached target")
