## 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 os
import sys
import math
import omni.kit.app
from omni.kit.test.async_unittest import AsyncTestCase
import omni.usd


class HardSoftRangeArrayUI(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        from omni.kit import ui_test
        from omni.kit.test_suite.helpers import open_stage, get_test_data_path, select_prims, wait_stage_loading, arrange_windows

        await arrange_windows("Stage", 64, 800)

        await open_stage(get_test_data_path(__name__, "usd/soft_range_float3.usda"))
        await ui_test.find("Property").focus()

        # select prim
        await select_prims(["/World/Looks/float3_softrange/Shader"])

        # wait for material to load & UI to refresh
        await wait_stage_loading()

    # After running each test
    async def tearDown(self):
        from omni.kit import ui_test
        from omni.kit.test_suite.helpers import select_prims, wait_stage_loading

        # de-select prims to prevent _delayed_dirty_handler exception
        await select_prims([])
        await wait_stage_loading()

    def _verify_values(self, widget_name: str, expected_value, model_soft_min, model_soft_max, hard_min, hard_max):
        from omni.kit import ui_test

        frame = ui_test.find(f"Property//Frame/**/.identifier=='{widget_name}'")
        for widget in frame.find_all(f"**/FloatDrag[*]"):
            self.assertAlmostEqual(widget.model.get_value_as_float(), expected_value.pop(0))
            if model_soft_min and widget.model._soft_range_min:
                self.assertAlmostEqual(widget.model._soft_range_min, model_soft_min)
            if model_soft_max and widget.model._soft_range_max:
                self.assertAlmostEqual(widget.model._soft_range_max, model_soft_max)

            self.assertAlmostEqual(widget.widget.min, hard_min)
            self.assertAlmostEqual(widget.widget.max, hard_max)

    async def _click_type(self, widget_name: str, values):
        from omni.kit import ui_test

        frame = ui_test.find(f"Property//Frame/**/.identifier=='{widget_name}'")
        for widget in frame.find_all(f"**/FloatDrag[*]"):
            await widget.input(str(values.pop(0)))


    async def test_hard_softrange_float3_ui(self):
        from omni.kit import ui_test

        # check default values
        self._verify_values("drag_per_channel_inputs:float3_1",
                      expected_value=[0.75, 0.5, 0.25],
                      model_soft_min=None,
                      model_soft_max=None,
                      hard_min=0,
                      hard_max=1)

        # change values
        await self._click_type("drag_per_channel_inputs:float3_1", values=[0.25, 1.45, 1.35])

        # verify changed values
        self._verify_values("drag_per_channel_inputs:float3_1",
                      expected_value=[0.25, 1.45, 1.35],
                      model_soft_min=0.0,
                      model_soft_max=1.25,
                      hard_min=0,
                      hard_max=1)

        # change values
        await self._click_type("drag_per_channel_inputs:float3_1", values=[23.25, 21.45, 24.5, 99.35])

        # verify changed values
        self._verify_values("drag_per_channel_inputs:float3_1",
                      [10.0, 10.0, 10.0, 10.0],
                      model_soft_min=0.0,
                      model_soft_max=10.0,
                      hard_min=0,
                      hard_max=10.0)

        # click reset
        frame = ui_test.find("Property//Frame/**/.identifier=='drag_per_channel_inputs:float3_1'")
        for widget in frame.find_all(f"**/ImageWithProvider[*]"):
            await widget.click()

        # check default values
        self._verify_values("drag_per_channel_inputs:float3_1",
                      expected_value=[0.75, 0.5, 0.25],
                      model_soft_min=None,
                      model_soft_max=None,
                      hard_min=0,
                      hard_max=1)


    async def test_hard_softrange_float4_ui(self):
        from omni.kit import ui_test

        # check default values
        self._verify_values("drag_per_channel_inputs:float4_1",
                      expected_value=[0.75, 0.5, 0.45, 0.25],
                      model_soft_min=None,
                      model_soft_max=None,
                      hard_min=0,
                      hard_max=1)

        # change values
        await self._click_type("drag_per_channel_inputs:float4_1", values=[1.25, 1.45, 1.5, 1.35])

        # verify changed values
        self._verify_values("drag_per_channel_inputs:float4_1",
                      expected_value=[1.25, 1.45, 1.5, 1.35],
                      model_soft_min=0.0,
                      model_soft_max=1.25,
                      hard_min=0,
                      hard_max=1.25)

        # change values
        await self._click_type("drag_per_channel_inputs:float4_1", values=[10.25, 10.45, 10.5, 10.35])

        # verify changed values
        self._verify_values("drag_per_channel_inputs:float4_1",
                      [5.0, 5.0, 5.0, 5.0],
                      model_soft_min=0.0,
                      model_soft_max=5.0,
                      hard_min=0,
                      hard_max=5.0)

        # click reset
        frame = ui_test.find("Property//Frame/**/.identifier=='drag_per_channel_inputs:float4_1'")
        for widget in frame.find_all(f"**/ImageWithProvider[*]"):
            await widget.click()

        # check default values
        self._verify_values("drag_per_channel_inputs:float4_1",
                      expected_value=[0.75, 0.5, 0.45, 0.25],
                      model_soft_min=None,
                      model_soft_max=None,
                      hard_min=0,
                      hard_max=1)


    async def test_hard_softrange_float2_ui(self):
        from omni.kit import ui_test

        # check default values
        self._verify_values("drag_per_channel_inputs:float2_1", [0.75, 0.25], None, None, 0, 1)

        # change values
        await self._click_type("drag_per_channel_inputs:float2_1", values=[0.25, 1.35])

        # verify changed values
        self._verify_values("drag_per_channel_inputs:float2_1",
                      expected_value=[0.25, 1.0],
                      model_soft_min=0.0,
                      model_soft_max=1.0,
                      hard_min=0,
                      hard_max=1)

        # change values
        await self._click_type("drag_per_channel_inputs:float2_1", values=[-2.45, -2.35])

        # verify changed values
        self._verify_values("drag_per_channel_inputs:float2_1",
                      [0.0, 0.0],
                      model_soft_min=0.0,
                      model_soft_max=1.0,
                      hard_min=0,
                      hard_max=1.0)

        # click reset
        frame = ui_test.find("Property//Frame/**/.identifier=='drag_per_channel_inputs:float2_1'")
        for widget in frame.find_all(f"**/ImageWithProvider[*]"):
            await widget.click()

        # check default values
        self._verify_values("drag_per_channel_inputs:float2_1", [0.75, 0.25], None, None, 0, 1)
