# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import asyncio
import os
import tempfile
import uuid
import carb
import omni.kit.commands
import omni.kit.test
import omni.kit.viewport_legacy
from pxr import Gf, Usd, UsdGeom, Sdf
from omni.kit.test_suite.helpers import arrange_windows


class TestViewportWindowCamera(omni.kit.test.AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows()

    # After running each test
    async def tearDown(self):
        pass

    async def test_builtin_camera_settings(self):
        async def test_builtin_camera_settings_inner(modify_camera=False):
            await omni.usd.get_context().new_stage_async()
            # Waiting one frame after the new stage, so in sync load cases
            # the USD context had a chance to launch a preSync event
            await omni.kit.app.get_app().next_update_async()
            viewport_window = omni.kit.viewport_legacy.get_default_viewport_window()
            omni.kit.commands.execute("CreatePrim", prim_type="Sphere")

            camera_paths = ["/OmniverseKit_Persp", "/OmniverseKit_Front", "/OmniverseKit_Top", "/OmniverseKit_Right"]
            camera_names = ["Perspective", "Front", "Top", "Right"]

            positions = []
            targets = []

            if modify_camera:
                viewport_window.set_camera_position("/OmniverseKit_Persp", 100, 200, 300, True)
                viewport_window.set_camera_target("/OmniverseKit_Persp", 300, 200, 100, True)
                viewport_window.set_active_camera("/OmniverseKit_Front")

            for i in range(len(camera_paths)):
                position = viewport_window.get_camera_position(camera_paths[i])
                self.assertTrue(position[0])
                positions.append(Gf.Vec3d(position[1], position[2], position[3]))

                target = viewport_window.get_camera_target(camera_paths[i])
                self.assertTrue(target[0])
                targets.append(Gf.Vec3d(target[1], target[2], target[3]))

            active_camera = viewport_window.get_active_camera()

            with tempfile.TemporaryDirectory() as tmpdirname:
                tmpfilename = os.path.join(tmpdirname, f"camTest_{str(uuid.uuid4())}.usda")
                carb.log_info("Saving temp camera path as %s" % (tmpfilename))
                await omni.usd.get_context().save_as_stage_async(tmpfilename)
                await omni.usd.get_context().reopen_stage_async()
                await omni.kit.app.get_app().next_update_async()
                for i in range(len(camera_names)):
                    carb.log_info("Testing %s" % (camera_names[i]))
                    # Position only saved for Persp camera
                    if i == 0:
                        saved_position = viewport_window.get_camera_position(camera_paths[i])
                        self.assertTrue(saved_position[0])
                        saved_position = Gf.Vec3d(saved_position[1], saved_position[2], saved_position[3])

                        carb.log_info("Comparing Positions: (%f,%f,%f) to (%f,%f,%f)" % (positions[i][0], positions[i][1], positions[i][2], saved_position[0], saved_position[1], saved_position[2]))
                        self.assertTrue(
                            Gf.IsClose(positions[i], saved_position, 0.000001),
                            f"Camera position mismatched! {positions[i]} is not close to {saved_position}",
                        )

                    saved_target = viewport_window.get_camera_target(camera_paths[i])
                    self.assertTrue(saved_target[0])
                    saved_target = Gf.Vec3d(saved_target[1], saved_target[2], saved_target[3])

                    active_camera_saved = viewport_window.get_active_camera()

                    carb.log_info("Comparing Targets: (%f,%f,%f) to (%f,%f,%f)" % (targets[i][0], targets[i][1], targets[i][2], saved_target[0], saved_target[1], saved_target[2]))
                    self.assertTrue(
                        Gf.IsClose(targets[i], saved_target, 0.000001),
                        f"Camera target mismatched! {targets[i]} is not close to {saved_target}",
                    )
                    if modify_camera and i == 0:
                        self.assertTrue(Gf.IsClose(positions[i], Gf.Vec3d(100, 200, 300), 0.000001))
                        self.assertTrue(Gf.IsClose(targets[i], Gf.Vec3d(300, 200, 100), 0.000001))

                    self.assertTrue(
                        active_camera == active_camera_saved,
                        f"Active camera mismatched! {active_camera} is not equal to {active_camera_saved}",
                    )

        # test default camera settings
        carb.log_info("Testing default camera settings")
        await test_builtin_camera_settings_inner()

        # test modified camera settings
        carb.log_info("Testing modified camera settings")
        await test_builtin_camera_settings_inner(True)

    async def test_camera_duplicate_from_static_active_camera(self):
        # Test for OM-35156
        await omni.usd.get_context().new_stage_async()
        await omni.kit.app.get_app().next_update_async()

        stage = omni.usd.get_context().get_stage()
        prim = UsdGeom.Xform.Define(stage, "/World")
        stage.SetDefaultPrim(prim.GetPrim())
        default_prim = stage.GetDefaultPrim()
        UsdGeom.XformCommonAPI(default_prim).SetTranslate(Gf.Vec3d(100, 0, 0))

        perspective_camera = stage.GetPrimAtPath("/OmniverseKit_Persp")
        persp_camera_world_mtx = omni.usd.get_world_transform_matrix(perspective_camera)
        default_prim_world_mtx = omni.usd.get_world_transform_matrix(default_prim)

        omni.kit.commands.execute(
            "SetActiveViewportCamera", new_active_cam_path=perspective_camera.GetPath(), viewport_name=""
        )
        await omni.kit.app.get_app().next_update_async()
        omni.kit.commands.execute("DuplicateFromActiveViewportCamera")
        await omni.kit.app.get_app().next_update_async()

        # The name of it will be Camera
        new_camera_prim = stage.GetPrimAtPath(default_prim.GetPath().AppendElementString("Camera"))
        new_camera_local_mtx = omni.usd.get_local_transform_matrix(new_camera_prim)
        local_mtx = persp_camera_world_mtx * default_prim_world_mtx.GetInverse()
        self.assertTrue(Gf.IsClose(new_camera_local_mtx, local_mtx, 0.000001))

    async def test_camera_duplicate_from_animated_active_camera(self):
        # Test for OM-33581
        layer_with_animated_camera = """\
        #usda 1.0
        (
            customLayerData = {
            }
            defaultPrim = "World"
            endTimeCode = 619
            framesPerSecond = 24
            metersPerUnit = 0.01
            startTimeCode = 90
            timeCodesPerSecond = 24
            upAxis = "Y"
        )

        def Xform "World" (
            kind = "assembly"
        )
        {
            def Camera "SHOTCAM"
            {
                float2 clippingRange = (10, 10000)
                custom bool depthOfField = 0
                float focalLength = 35
                float focalLength.timeSamples = {
                    172: 35,
                    173: 34.650208
                }
                float focusDistance = 5
                float fStop = 5.6
                float horizontalAperture = 35.999928
                float verticalAperture = 20.22853
                float3 xformOp:rotateXYZ.timeSamples = {
                    90: (8.786562, -19.858515, -1.2905762),
                    91: (8.744475, -19.858515, -1.274945)
                }
                float3 xformOp:scale.timeSamples = {
                    90: (12.21508, 12.21508, 12.21508)
                }
                double3 xformOp:translate.timeSamples = {
                    90: (-462.44292428348126, 1734.9897515804505, -2997.4519412288555),
                    151: (-462.4429242834828, 1734.9897515804498, -2997.451941228857)
                }
                uniform token[] xformOpOrder = ["xformOp:translate", "xformOp:rotateXYZ", "xformOp:scale"]
            }
        }
        """

        format = Sdf.FileFormat.FindByExtension(".usd")
        root_layer = Sdf.Layer.CreateAnonymous("world.usd", format)
        root_layer.ImportFromString(layer_with_animated_camera)
        stage = Usd.Stage.Open(root_layer)
        await omni.usd.get_context().attach_stage_async(stage)
        await omni.kit.app.get_app().next_update_async()

        old_camera = stage.GetPrimAtPath("/World/SHOTCAM")
        omni.kit.commands.execute(
            "SetActiveViewportCamera", new_active_cam_path="/World/SHOTCAM", viewport_name=""
        )
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        omni.kit.commands.execute("DuplicateFromActiveViewportCamera")
        await omni.kit.app.get_app().next_update_async()

        new_camera = stage.GetPrimAtPath("/World/Camera")
        self.assertTrue(new_camera)

        timeline = omni.timeline.get_timeline_interface()
        timecode = timeline.get_current_time() * stage.GetTimeCodesPerSecond()
        old_matrix = omni.usd.get_world_transform_matrix(old_camera, timecode)
        new_matrix = omni.usd.get_world_transform_matrix(new_camera)
        self.assertTrue(Gf.IsClose(new_matrix, old_matrix, 0.000001))

        # Set time code to the second timecode of rotateXYZ, so they are not equal.
        timeline.set_current_time(float(91.0 / stage.GetTimeCodesPerSecond()))
        timecode = timeline.get_current_time() * stage.GetTimeCodesPerSecond()
        old_matrix = omni.usd.get_world_transform_matrix(old_camera, timecode)
        self.assertFalse(Gf.IsClose(new_matrix, old_matrix, 0.000001))

        # Make another duplicate that's duplicated from updated time.
        omni.kit.commands.execute(
            "SetActiveViewportCamera", new_active_cam_path="/World/SHOTCAM", viewport_name=""
        )
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        omni.kit.commands.execute("DuplicateFromActiveViewportCamera")
        await omni.kit.app.get_app().next_update_async()

        new_camera = stage.GetPrimAtPath("/World/Camera_01")
        self.assertTrue(new_camera)
        old_matrix = omni.usd.get_world_transform_matrix(old_camera, timecode)
        new_matrix = omni.usd.get_world_transform_matrix(new_camera)
        self.assertTrue(Gf.IsClose(new_matrix, old_matrix, 0.000001))

    async def test_layer_clear_to_remove_camera_deltas(self):
        # Test for OM-47804
        await omni.usd.get_context().new_stage_async()
        for _ in range(10):
            # Create anonymous layer and set edit target
            stage = omni.usd.get_context().get_stage()
            replicator_layer = None
            for layer in stage.GetLayerStack():
                if layer.GetDisplayName() == "Replicator":
                    replicator_layer = layer
                    break

            if replicator_layer is None:
                root = stage.GetRootLayer()
                replicator_layer = Sdf.Layer.CreateAnonymous(tag="Replicator")
                root.subLayerPaths.append(replicator_layer.identifier)
            omni.usd.set_edit_target_by_identifier(stage, replicator_layer.identifier)

            # Clear prim and assert not valid
            replicator_layer.Clear()
            self.assertFalse(stage.GetPrimAtPath("/ReplicatorPrim").IsValid())

            # Define some prim on the layer
            stage.DefinePrim("/ReplicatorPrim", "Camera")
