## 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.commands
import omni.kit.test
from .test_common import RtxTest, testSettings, postLoadTestSettings, wait_for_update
from pxr import UsdGeom


spherePath = '/World/Sphere'
cylinderPath = '/World/Cylinder'

defaultColor0 = [0, 0, 0, 0]
defaultColor1 = [0.7, 0.7, 0, 1]
defaultColor2 = [0.7, 0, 1, 1]

displayOptionSetting = "/persistent/app/viewport/displayOptions"
outlineWidthSetting = "/persistent/app/viewport/outline/width"
outlineColorSetting = "/persistent/app/viewport/outline/color"
shadeColorSetting = "/persistent/app/viewport/outline/shadeColor"
intersectionColorSetting = "/persistent/app/viewport/outline/intersection/color"
outlineEnabledSetting = "/app/viewport/outline/enabled"


class TestRtxSelectionOutline(RtxTest):
    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)
        self.set_settings({displayOptionSetting: (1 << 7)})
        omni.usd.get_context().new_stage()  # Important: it should be called before `setUp` method as it reset settings
                                            # and stage.
        await omni.kit.app.get_app().next_update_async()  # Wait stage loading
        self.set_settings(postLoadTestSettings)
        self.set_settings({
            outlineWidthSetting: 2,
            outlineColorSetting: defaultColor0,
            shadeColorSetting: defaultColor0,
            intersectionColorSetting: defaultColor0,
            "/app/transform/operation": "select"
        })

        self.create_scene()
        self.ctx.set_selection_group_outline_color(255, defaultColor1)

    def set_selected(self, path, select=True):
        self.ctx.get_selection().set_prim_path_selected(path, select, True, False, True)

    def create_scene(self):
        stage = self.ctx.get_stage()
        sphere = UsdGeom.Sphere.Define(stage, spherePath)
        sphere.CreateRadiusAttr(100)
        cylinder = UsdGeom.Cylinder.Define(stage, cylinderPath)
        cylinder.CreateHeightAttr(200)
        cylinder.CreateRadiusAttr(75)

    async def test_rtx_selectionoutline_outline1(self):
        """
        Test SelectionOutline - Outline sphere
        """
        self.set_selected(spherePath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Outline1.png")

    async def test_rtx_selectionoutline_outline2(self):
        """
        Test SelectionOutline - Outline cylinder
        """
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Outline2.png")

    async def test_rtx_selectionoutline_outline12(self):
        """
        Test SelectionOutline - Outline sphere & cylinder
        """
        self.set_selected(spherePath)
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Outline12.png")

    async def test_rtx_selectionoutline_intersection(self):
        """
        Test SelectionOutline - Intersection
        """
        self.set_settings({intersectionColorSetting: defaultColor2})
        self.set_selected(spherePath)
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Intersection.png", 4e-5)

    async def test_UNSTABLE_rtx_selectionoutline_thick(self):
        """
        Test SelectionOutline - Outline sphere & cylinder, thick
        """
        self.set_settings({outlineWidthSetting: 15})
        self.set_selected(spherePath)
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Thick.png")

    async def test_rtx_selectionoutline_blend(self):
        """
        Test SelectionOutline - Outline cylinder with blended color
        """
        self.set_settings({outlineWidthSetting: 15})
        c = defaultColor1[:3]
        c.append(0.5)
        self.ctx.set_selection_group_outline_color(255, c)
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Blend.png")

    async def test_rtx_selectionoutline_multicolored(self):
        """
        Test SelectionOutline - Outline sphere & cylinder with different colors
        """
        self.set_settings({outlineWidthSetting: 4})
        self.ctx.set_selection_group(1, spherePath)
        self.ctx.set_selection_group(2, cylinderPath)
        self.ctx.set_selection_group_outline_color(1, defaultColor1)
        self.ctx.set_selection_group_outline_color(2, defaultColor2)
        self.ctx.set_selection_group_shade_color(1, defaultColor0)
        self.ctx.set_selection_group_shade_color(2, defaultColor0)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Multicolored.png")

    async def test_rtx_selectionoutline_shade(self):
        """
        Test SelectionOutline - Shade sphere
        """
        self.set_settings({outlineWidthSetting: 4})
        self.ctx.set_selection_group(1, spherePath)
        self.ctx.set_selection_group(2, cylinderPath)
        self.ctx.set_selection_group_shade_color(1, defaultColor1)
        self.ctx.set_selection_group_shade_color(2, defaultColor0)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Shade.png")

    async def test_UNSTABLE_rtx_selectionoutline_shade_multicolored_UNSTABLE(self):
        """
        Test SelectionOutline - Shade sphere & cylinder with different colors
        """
        self.set_settings({outlineWidthSetting: 4})
        self.ctx.set_selection_group(1, spherePath)
        self.ctx.set_selection_group(2, cylinderPath)
        self.ctx.set_selection_group_shade_color(1, defaultColor1)
        self.ctx.set_selection_group_shade_color(2, defaultColor2)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "ShadeMulticolored.png")

    async def test_rtx_selectionoutline_disabled(self):
        """
        Test SelectionOutline - Disabled
        """
        self.set_settings({outlineEnabledSetting: False, displayOptionSetting: 0})
        self.set_selected(spherePath)
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Disabled.png")

    # OM-102267 - unreliable in 105.1 branch
    async def test_UNSTABLE_rtx_selectionoutline_disabledByMetadata_UNSTABLE(self):
        """
        Test SelectionOutline - Disabled by metadata.
        OM-34048
        """
        self.ctx.get_stage().GetPrimAtPath(spherePath).SetMetadata('no_selection_outline', True)
        self.set_selected(spherePath)
        self.set_selected(cylinderPath)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "Outline2.png", 2e-5)

    # OM-102267 - unreliable in 105.1 branch
    async def test_UNSTABLE_rtx_selectionoutline_instances_UNSTABLE(self):
        """
        Test Selection Outline - Instances
        """
        self.open_usd("hydra/CubeWorld.usda")
        PRIM_PATH = '/CubeWorld/CubeInstance'
        cubeInstance = self.ctx.get_stage().GetPrimAtPath(PRIM_PATH)
        cubeInstance.SetInstanceable(False)
        self.set_selected(PRIM_PATH)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "InstanceableFalse.png", 4e-5)
        self.set_selected(PRIM_PATH, False) # Unselect

        cubeInstance.SetInstanceable(True)
        self.set_selected(PRIM_PATH)
        await wait_for_update()
        await self.capture_and_compare("SelectionOutline", "InstanceableTrue.png", 4e-5)

    async def test_rtx_selectionoutline_mouse_select(self):
        """Test mouse selection works in Viewport."""
        from omni.kit import ui_test
        from omni.kit.ui_test.vec2 import Vec2

        self.open_usd("hydra/CubeInstance.usda")
        omni.usd.get_context().get_selection().set_selected_prim_paths([], True)
        await wait_for_update()

        await ui_test.emulate_mouse_move_and_click(Vec2(256, 256))
        await wait_for_update()

        await self.capture_and_compare("SelectionOutline", "MouseSelect.png", 4e-5)

    async def test_rtx_selectionoutline_mouse_select_aov_changed(self):
        """Test mouse selection works in Viewport before and after AOV change."""
        from omni.kit import ui_test
        from omni.kit.ui_test.vec2 import Vec2

        mouse_pos = Vec2(256, 256)
        self.open_usd("hydra/CubeInstance.usda")
        usd_context = omni.usd.get_context()
        usd_context.get_selection().set_selected_prim_paths([], True)
        await wait_for_update()

        await ui_test.emulate_mouse_move_and_click(mouse_pos)
        await wait_for_update()

        await self.capture_and_compare("SelectionOutline", "MouseSelectAOVA.png", 4e-5)
        # Reset selection to nothing
        usd_context.get_selection().set_selected_prim_paths([], True)

        # Add additional AOVs
        from omni.kit.viewport.utility import get_active_viewport
        from pxr import Sdf, Usd, UsdRender
        stage = usd_context.get_stage()
        render_product_path = Sdf.Path(get_active_viewport().render_product_path)
        render_product = UsdRender.Product(stage.GetPrimAtPath(render_product_path))
        ordered_vars_rel = render_product.GetOrderedVarsRel()
        start_aov_len = len(ordered_vars_rel.GetForwardedTargets())

        # Targets need to be adde don session-layer as other prims live there
        with Usd.EditContext(stage, stage.GetSessionLayer()):
            aov_name = "Depth"
            render_var = UsdRender.Var.Define(stage, Sdf.Path(f"/Render/Vars/{aov_name}"))
            render_var.GetSourceNameAttr().Set(aov_name)
            render_var.GetDataTypeAttr().Set("float")
            ordered_vars_rel.AddTarget(render_var.GetPath())

        end_aov_len = len(render_product.GetOrderedVarsRel().GetForwardedTargets())
        self.assertNotEqual(start_aov_len, end_aov_len)

        # Re-select the object
        await ui_test.emulate_mouse_move_and_click(mouse_pos)
        await wait_for_update()

        await self.capture_and_compare("SelectionOutline", "MouseSelectAOVB.png", 4e-5)
