## 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.test
import omni.timeline
from .test_hydra_common import RtxHydraTest, RtxHydraMaterialsTest, RtxHydraInstancerTest
from .test_common import wait_for_update, postLoadTestSettings
from pxr import Sdf, UsdGeom, UsdShade

DEFAULT_POINTS = [(-50, -50, -50), (50, -50, -50), (-50, -50, 50), (50, -50, 50), (-50, 50, -50), (50, 50, -50), (50, 50, 50), (-50, 50, 50)]

class TestRtxHydraPoints(RtxHydraTest):
    TEST_PATH = "hydra/points"
    PRIM_PATH = "/World/pts"

    async def display_color(self):
        # Constant
        geom = self.create_geometry()
        geom.CreateDisplayColorPrimvar("constant").Set([(0, 1, 0)])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "displayColorConstant.png")
        return geom

    async def widths(self):
        # Default
        geom = self.create_geometry()
        geom.GetWidthsAttr().Clear()
        self.set_settings({"/persistent/rtx/hydra/points/defaultWidth": 35})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "widthDefault.png")
        self.set_settings({"/persistent/rtx/hydra/points/defaultWidth": 45})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "widthDefault2.png")
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        # Constant
        geom = self.create_geometry()
        geom.SetWidthsInterpolation("constant")
        geom.CreateWidthsAttr().Set([20])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "widthConstant.png")
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
        # Vertex
        geom = self.create_geometry()
        geom.CreateWidthsAttr().Set([20, 15, 10, 25, 30, 35, 40, 45])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "widthVertex.png")

    async def points(self):
        pts = [DEFAULT_POINTS,
               [(-50, -50, -50), (150, -50, -50), (-50, -50, 50), (50, -50, 50), (-50, 50, -50), (50, 50, -50), (50, 50, 50), (-50, 50, 50)]]
        await self.attribute_test_case(self.create_geometry, "points", Sdf.ValueTypeNames.Point3f, pts, ["geom.png", "points.png"])

    async def skipProcessing(self):
        geom = self.create_geometry()
        geom.GetPrim().CreateAttribute("omni:rtx:skip", Sdf.ValueTypeNames.Bool).Set(True)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "empty.png")

class TestRtxHydraPrimitiveDrawingPoints(TestRtxHydraPoints):
    TEST_PATH = "hydra/points/primitivedrawing"
    PRIM_PATH = "/World/pts"

    def create_geometry(self, name=PRIM_PATH):
        pts = UsdGeom.Points.Define(self.ctx.get_stage(), name)
        pts.CreatePointsAttr(DEFAULT_POINTS)
        pts.CreateWidthsAttr([10, 10, 10, 10, 5, 5, 15, 15])
        UsdGeom.PrimvarsAPI(pts).CreatePrimvar("usePrimitiveDrawing", Sdf.ValueTypeNames.Bool).Set(True)
        UsdGeom.PrimvarsAPI(pts).CreatePrimvar("screenSpacePrimitiveDrawing", Sdf.ValueTypeNames.Bool).Set(True)
        return pts

    async def test_visibility(self):
        """
        Test hydra points - visibility
        """
        await self.visibility()

    async def test_transform(self):
        """
        Test hydra points - tranform
        """
        await self.transform_all()

    async def test_displayColor(self):
        """
        Test hydra points - change display color
        """
        geom = await self.display_color()

        # Vertex
        geom.CreateDisplayColorPrimvar("vertex").Set([(0, 1, 0), (1, 1, 0), (0, 1, 1), (1, 0, 1), (1, 0, 0), (0, 0, 1), (0, 1, 0), (1, 1, 0)])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "displayColorVertex.png")

    async def test_widths(self):
        """
        Test hydra points - change widths
        """
        await self.widths()

    async def test_points(self):
        """
        Test hydra points - update points
        """
        await self.points()

    async def test_skipProcessing(self):
        """
        Test hydra points - skipProcessing attribute
        """
        await self.skipProcessing()


class TestRtxHydraPrimitiveDrawingPointsInstancer(RtxHydraInstancerTest):
    # TODO full cover scene instancer with tests
    TEST_PATH = "hydra/points/primitivedrawing/instancer"
    GEOM_PATH = RtxHydraInstancerTest.PRIM_PATH + "/pts"

    def create_geometry(self, name=GEOM_PATH):
        pts = UsdGeom.Points.Define(self.ctx.get_stage(), name)
        pts.CreatePointsAttr(DEFAULT_POINTS)
        pts.CreateWidthsAttr([10, 10, 10, 10, 5, 5, 15, 15])
        UsdGeom.PrimvarsAPI(pts).CreatePrimvar("usePrimitiveDrawing", Sdf.ValueTypeNames.Bool).Set(True)
        UsdGeom.PrimvarsAPI(pts).CreatePrimvar("screenSpacePrimitiveDrawing", Sdf.ValueTypeNames.Bool).Set(True)
        return pts

    async def test_sceneInstancer(self):
        """
        Test hydra points scene instancer
        """
        await self.si_all()

    async def test_pointInstancer(self):
        """
        Test hydra points point instancer
        """
        await self.pi_all()


class UNSTABLE_TestRtxHydraProceduralPoints_UNSTABLE(TestRtxHydraPoints):
    TEST_PATH = "hydra/points/procedural"
    PRIM_PATH = "/World/pts"

    def create_geometry(self, name=PRIM_PATH):
        pts = UsdGeom.Points.Define(self.ctx.get_stage(), name)
        pts.CreatePointsAttr(DEFAULT_POINTS)
        pts.CreateWidthsAttr([40, 30, 30, 20, 15, 15, 35, 35])
        return pts

    async def test_visibility(self):
        """
        Test hydra points - visibility
        """
        await self.visibility()

    async def test_doNotCastShadow(self):
        """
        Test hydra points - doNotCastShadow toggle
        """
        self.create_geometry_with_floor()
        await self.toggle_primvar(self.PRIM_PATH, "doNotCastShadows")

    async def test_matteObject(self):
        """
        Test hydra points - isMatteObject toggle
        """
        self.create_geometry_with_floor()
        await self.matte()

    async def test_hideForCamera(self):
        """
        Test hydra points - hideForCamera toggle
        """
        self.create_geometry_with_floor()
        await self.toggle_primvar(self.PRIM_PATH, "hideForCamera")

    async def test_transform(self):
        """
        Test hydra points - tranform
        """
        await self.transform_all()

    async def test_displayColor(self):
        """
        Test hydra points - change display color
        """
        await self.display_color()

    async def test_widths(self):
        """
        Test hydra points - change widths
        """
        await self.widths()

    async def test_points(self):
        """
        Test hydra points - update points
        """
        await self.points()
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)

        # OM-45604, OM-43655
        # Delayed points assignment
        geom = UsdGeom.Points.Define(self.ctx.get_stage(), self.PRIM_PATH)
        geom.CreateWidthsAttr([40, 30, 30, 20, 15, 15, 35, 35])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "empty.png", 1e-4)

        attr = geom.CreatePointsAttr(DEFAULT_POINTS)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "geom.png", 1e-4)

        # OM-43655 - remove point
        attr.Set(DEFAULT_POINTS[:-1])
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "pointsRemovedOne.png", 1e-4)

    async def test_skip_processing(self):
        """
        Test hydra points - skipProcessing attribute
        """
        await self.skipProcessing()

    async def test_texCoords(self):
        self.set_settings({"/rtx/debugView/target": "texcoord1"})
        self.create_geometry()
        prim = self.ctx.get_stage().GetPrimAtPath(self.PRIM_PATH)
        attr = UsdGeom.PrimvarsAPI(prim).CreatePrimvar("uv1", Sdf.ValueTypeNames.Float2Array, UsdGeom.Tokens.vertex)
        uvVertexSet = [(1, 0), (0.15, 0.5), (0.3, 0.5), (0.5, 0.5), (0.7, 0.5), (0.85, 0.5), (1, 0.5), (1, 1)]
        attr.Set(uvVertexSet)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "texCoords.png")
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)

    async def test_timeSampled(self):
        """
        Test hydra points - verify time sample animation works.
        """
        super().open_usd("hydra/points_timeSampled_001.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

        timeline = omni.timeline.get_timeline_interface()

        # Regression test for OM-96462, where time-sampled xform was impairing update of time-sampled points
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "points_timeSampled_frame0.png")

        # Advance a few frames and capture another frame
        for t in range(12):
            timeline.forward_one_frame()
            await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "points_timeSampled_frame12.png")


class UNSTABLE_TestRtxHydraProceduralPointsInstancer_UNSTABLE(RtxHydraInstancerTest):
    # TODO full cover scene instancer with tests
    TEST_PATH = "hydra/points/procedural/instancer"
    GEOM_PATH = RtxHydraInstancerTest.PRIM_PATH + "/pts"

    def create_geometry(self, name=GEOM_PATH):
        pts = UsdGeom.Points.Define(self.ctx.get_stage(), name)
        pts.CreatePointsAttr(DEFAULT_POINTS)
        pts.CreateWidthsAttr([40, 30, 30, 20, 15, 15, 35, 35])
        return pts

    async def test_sceneInstancer(self):
        """
        Test hydra points - scene instancer
        """
        await self.si_all()

    async def test_pointInstancer(self):
        """
        Test hydra points - point instancer
        """
        await self.pi_all()


class UNSTABLE_TestRtxHydraProceduralPointsMaterials_UNSTABLE(RtxHydraMaterialsTest):
    TEST_PATH = "hydra/points/procedural/material"
    PRIM_PATH = "/World/pts"

    def create_geometry(self, name=PRIM_PATH):
        pts = UsdGeom.Points.Define(self.ctx.get_stage(), name)
        pts.CreatePointsAttr([(-50, 0, 0), (50, 0, 0)])
        pts.CreateWidthsAttr([50, 25])
        return pts

    async def test_materials(self):
        """
        Test hydra points - materials
        """
        materials = ["Green", "Red"]
        looksPath = "/World/Looks/"
        # Set and change material
        geomPrim = self.create_geometry().GetPrim()
        for m in materials:
            self.bind_material(geomPrim, looksPath + m)
            await wait_for_update()
            await self.capture_and_compare(self.TEST_PATH, "mat_{}.png".format(m))

        # Unbind material it should fallback to basic material
        UsdShade.MaterialBindingAPI(geomPrim).UnbindAllBindings()
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "mat_base.png")
        self.ctx.get_stage().RemovePrim(self.PRIM_PATH)
