import traceback
import uuid
import carb
import carb.settings
import omni.ext
import omni.ui as ui
import omni.gpu_foundation_factory
import omni.client
from random import random
import math
from omni.physx.scripts import physicsUtils
from pxr import UsdLux, UsdGeom, Gf, UsdPhysics


def print_exception(exc_info):
    print(f"===== Exception: {exc_info[0]}, {exc_info[1]}. Begin Stack Trace =====")
    traceback.print_tb(exc_info[2])
    print("===== End Stack Trace ===== ")

class Extension(omni.ext.IExt):
    def on_startup(self):
        # Obtain multi process info
        gpu = omni.gpu_foundation_factory.get_gpu_foundation_factory_interface()
        processIndex = gpu.get_process_index()
        processCount = gpu.get_process_count()

        # Settings
        self._settings = carb.settings.get_settings()

        self._settings.set("/persistent/physics/updateToUsd", False)
        self._settings.set("/persistent/physics/updateVelocitiesToUsd", False)
        self._settings.set("/persistent/omnihydra/useFastSceneDelegate", True)
        self._settings.set("/app/renderer/sleepMsOutOfFocus", 0)

        if processIndex > 0:
            self._settings.set("/app/window/title", "Kit Worker {} //".format(processIndex))

        # UI
        self._window = ui.Window("RTX Multi-Node Prototype", dockPreference=ui.DockPreference.LEFT_BOTTOM)
        self._window.deferred_dock_in("Console", ui.DockPolicy.CURRENT_WINDOW_IS_ACTIVE)

        with self._window.frame:
            with omni.ui.HStack(height=0):
                omni.ui.Spacer(width=6)
                with omni.ui.VStack(height=0):
                    omni.ui.Spacer(height=10)

                    with omni.ui.HStack(height=0):
                        if processIndex == 0:
                            omni.ui.Label("Master process of {}".format(processCount))
                        else:
                            omni.ui.Label("Worker process {} of {}".format(processIndex, processCount))

                    if processIndex == 0:
                        omni.ui.Spacer(height=20)

                        with omni.ui.HStack():
                            with omni.ui.VStack(height=0):
                                omni.ui.Label("Settings")
                                omni.ui.Spacer(height=5)

                                with omni.ui.HStack():
                                    omni.ui.Spacer(width=10)

                                    with omni.ui.VStack(width = 200, height = 0):
                                        with omni.ui.HStack():
                                            omni.ui.Label("Fabric Replication")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/flatcacheReplication/enabled"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/flatcacheReplication/enabled", a.get_value_as_bool()))
                                            omni.ui.Label("Enabled")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/flatcacheReplication/profile"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/flatcacheReplication/profile", a.get_value_as_bool()))
                                            omni.ui.Label("Profile")

                                        omni.ui.Spacer(height=10)

                                        with omni.ui.HStack():
                                            omni.ui.Label("Gather Render Results")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/gatherResults/enabled"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/gatherResults/enabled", a.get_value_as_bool()))
                                            omni.ui.Label("Enabled")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/gatherResults/profile"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/gatherResults/profile", a.get_value_as_bool()))
                                            omni.ui.Label("Profile")

                                    with omni.ui.VStack(width = 200, height=0):
                                        with omni.ui.HStack():
                                            omni.ui.Label("Render Product Replication")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/renderProductReplication/enabled"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/renderProductReplication/enabled", a.get_value_as_bool()))
                                            omni.ui.Label("Enabled")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/renderProductReplication/includeCameraAttributes"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/renderProductReplication/includeCameraAttributes", a.get_value_as_bool()))
                                            omni.ui.Label("Include Camera Attributes")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/renderProductReplication/fullSessionLayer"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/renderProductReplication/fullSessionLayer", a.get_value_as_bool()))
                                            omni.ui.Label("Full Session Layer")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/renderProductReplication/profile"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/renderProductReplication/profile", a.get_value_as_bool()))
                                            omni.ui.Label("Profile")
                                        omni.ui.Spacer(height=5)
                                        with omni.ui.HStack():
                                            omni.ui.Spacer(width=20)
                                            box = omni.ui.CheckBox(width = 22)
                                            box.model.set_value(self._settings.get("/multiNode/renderProductReplication/print"))
                                            box.model.add_value_changed_fn(lambda a: self._settings.set("/multiNode/renderProductReplication/print", a.get_value_as_bool()))
                                            omni.ui.Label("Print")

                            omni.ui.Spacer(width=200)

                            with omni.ui.VStack(height=0):
                                omni.ui.Label("Deprecated: Physics prototype")
                                omni.ui.Spacer(height=10)

                                with omni.ui.HStack(height=0):
                                    load_button = omni.ui.Button("Load Stage", width=150, height=30)
                                    load_button.set_clicked_fn(self._on_load_button_fn)

                                omni.ui.Spacer(height=5)

                                with omni.ui.HStack(height=0):
                                    self._numBoxesDrag = omni.ui.IntDrag(min = 1, max = 1000000, step = 1, width = 100)
                                    self._numBoxesDrag.model.set_value(1000)
                                    omni.ui.Spacer(width=10)
                                    omni.ui.Label("Boxes")

                                omni.ui.Spacer(height=5)

                                with omni.ui.HStack(height=0):
                                    init_button = omni.ui.Button("Init Simulation", width=150, height=30)
                                    init_button.set_clicked_fn(self._on_init_button_fn)

                                omni.ui.Spacer(height=5)

                                with omni.ui.HStack(height=0):
                                    play_button = omni.ui.Button("Play Simulation", width=150, height=30)
                                    play_button.set_clicked_fn(self._on_play_button_fn)

                    omni.ui.Spacer()

    def on_shutdown(self):
        pass

    def _on_load_button_fn(self):
        carb.log_warn("Load stage")

        srcPath = "omniverse://content.ov.nvidia.com/Users/tbiedert@nvidia.com/multi-node/empty.usd"
        dstPath = "omniverse://content.ov.nvidia.com/Users/tbiedert@nvidia.com/multi-node/gen/" + str(uuid.uuid4()) + ".usd"

        # Create a temporary copy
        try:
            result = omni.client.copy(srcPath, dstPath, omni.client.CopyBehavior.OVERWRITE)
            if result != omni.client.Result.OK:
                carb.log_error(f"Cannot copy from {srcPath} to {dstPath}, error code: {result}.")
            else:
                carb.log_warn(f"Copied from {srcPath} to {dstPath}.")
        except Exception as e:
            traceback.print_exc()
            carb.log_error(str(e))

        # Load the temporary copy
        omni.usd.get_context().open_stage(str(dstPath))

        # FIXME: Add compatible API for legacy usage.
        # Enable live sync
        # omni.usd.get_context().set_stage_live(layers.LayerLiveMode.ALWAYS_ON)

    def _on_init_button_fn(self):
        carb.log_warn("Init simulation")

        # Disable FlatCache replication
        self._settings.set("/multiNode/flatcacheReplication/enabled", False)

        # set up axis to z
        stage = omni.usd.get_context().get_stage()
        UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
        UsdGeom.SetStageMetersPerUnit(stage, 0.01)

        # Disable default light
        defaultPrimPath = str(stage.GetDefaultPrim().GetPath())
        stage.RemovePrim(defaultPrimPath + "/defaultLight")

        # light
        sphereLight = UsdLux.SphereLight.Define(stage, defaultPrimPath + "/SphereLight")
        sphereLight.CreateRadiusAttr(30)
        sphereLight.CreateIntensityAttr(50000)
        sphereLight.AddTranslateOp().Set(Gf.Vec3f(0.0, 0.0, 150.0))

        # Boxes
        numBoxes = self._numBoxesDrag.model.get_value_as_int()
        size = 25.0
        radius = 100.0

        self._boxes = [None] * numBoxes
        for i in range(numBoxes):
            box = UsdGeom.Cube.Define(stage, defaultPrimPath + "/box_" + str(i))
            box.CreateSizeAttr(size)
            box.CreateExtentAttr([(-size / 2, -size / 2, -size / 2), (size / 2, size / 2, size / 2)])
            box.AddTranslateOp().Set(Gf.Vec3f(radius * math.cos(0.5 * i), radius * math.sin(0.5 * i), size + (i * 0.2) * size))
            box.AddOrientOp().Set(Gf.Quatf(random(), random(), random(), random()).GetNormalized())
            box.CreateDisplayColorAttr().Set([Gf.Vec3f(random(), random(), random())])
            self._boxes[i] = box.GetPrim()

        # Plane
        planeActorPath = defaultPrimPath + "/plane"
        size = 1500.0
        planeGeom = UsdGeom.Cube.Define(stage, planeActorPath)
        self._plane = stage.GetPrimAtPath(planeActorPath)
        planeGeom.CreateSizeAttr(size)
        planeGeom.CreateExtentAttr([(-size / 2, -size / 2, -size / 2), (size / 2, size / 2, size / 2)])
        planeGeom.AddTranslateOp().Set(Gf.Vec3f(0.0, 0.0, 0.0))
        planeGeom.AddOrientOp().Set(Gf.Quatf(1.0))
        planeGeom.AddScaleOp().Set(Gf.Vec3f(1.0, 1.0, 0.01))
        planeGeom.CreateDisplayColorAttr().Set([Gf.Vec3f(0.5)])

    def _on_play_button_fn(self):
        carb.log_warn("Play simulation")

        # FIXME: Add compatible API for legacy usage.
        # Disable live sync
        # omni.usd.get_context().set_stage_live(layers.LayerLiveMode.TOGGLE_OFF)

        # Enable FlatCache replication
        self._settings.set("/multiNode/flatcacheReplication/enabled", True)

        # Add physics to scene
        stage = omni.usd.get_context().get_stage()
        defaultPrimPath = str(stage.GetDefaultPrim().GetPath())
        scene = UsdPhysics.Scene.Define(stage, defaultPrimPath + "/physicsScene")
        scene.CreateGravityDirectionAttr().Set(Gf.Vec3f(0.0, 0.0, -1.0))
        scene.CreateGravityMagnitudeAttr().Set(981.0)

        for box in self._boxes:
            UsdPhysics.CollisionAPI.Apply(box)
            physicsAPI = UsdPhysics.RigidBodyAPI.Apply(box)
            UsdPhysics.MassAPI.Apply(box)

        UsdPhysics.CollisionAPI.Apply(self._plane)

        # Play animation
        omni.timeline.get_timeline_interface().play()
