## Copyright (c) 2021, 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 omni.kit.app
import omni.kit.test
from .test_common import RtxTest, testSettings, postLoadTestSettings, wait_for_update
from pxr import UsdGeom, Sdf
import numpy as np


class TestRtxSceneDb(RtxTest):
    TEST_PATH = "scenedb"

    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)
        omni.usd.get_context().new_stage()  # Important: it should be called before `setUp` method as it reset settings and stage.
        self.add_dir_light()
        await omni.kit.app.get_app().next_update_async()  # Wait stage loading
        self.set_settings(postLoadTestSettings)

    def create_geometry(self, name):
        sphere = UsdGeom.Sphere.Define(self.ctx.get_stage(), name)
        sphere.CreateRadiusAttr(10)

    def create_instancer(self, name, N):
        pts = []
        indices = []
        sp = 250.
        dx = 500. / N[0]
        dy = 500. / N[1]
        dz = 500. / N[2]
        for i in np.arange(-sp, sp, dx):
            for j in np.arange(-sp, sp, dy):
                for k in np.arange(-sp, sp, dz):
                    pts.append((i, j, k))
                    indices.append(0)
        print("Generated points: ", len(indices))
        lyr = self.ctx.get_stage().GetRootLayer()
        with Sdf.ChangeBlock():
            instancer = lyr.GetPrimAtPath("/World/Instancer") or \
               Sdf.PrimSpec(lyr.GetPrimAtPath("/World"), "Instancer", Sdf.SpecifierDef, "PointInstancer")
            pos = lyr.GetAttributeAtPath("/World/Instancer.positions") or \
                Sdf.AttributeSpec(instancer, "positions", Sdf.ValueTypeNames.Point3fArray)
            pos.default = pts
            protoIndices = lyr.GetAttributeAtPath("/World/Instancer.protoIndices") or \
                Sdf.AttributeSpec(instancer, "protoIndices", Sdf.ValueTypeNames.IntArray)
            protoIndices.default = indices
            rel = lyr.GetRelationshipAtPath("/World/Instancer.prototypes") or \
                Sdf.RelationshipSpec(instancer, "prototypes")
            rel.targetPathList.prependedItems = [name]

    async def test_sphere_PI_1M(self):
        id = "/World/Sphere"
        self.create_geometry(id)
        self.create_instancer(id, (100, 100, 100))
        await wait_for_update(self.ctx, 25)  # Wait 25 frames before capture.
        await self.capture_and_compare(self.TEST_PATH, "sphere1M.png", threshold=5e-5)

    # Unstable due to crashes in SceneDB: OM-49598
    async def test_UNSTABLE_sphere_PI_1M_Recreate_UNSTABLE(self):
        id = "/World/Sphere"
        self.create_geometry(id)
        self.create_instancer(id, (100, 100, 100))
        await wait_for_update(self.ctx, 25)  # Wait 25 frames before capture.
        self.ctx.get_stage().RemovePrim("/World/instancer")
        self.create_instancer(id, (100, 100, 100))
        await self.capture_and_compare(self.TEST_PATH, "sphere1M.png")
