
## 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 carb
import omni.kit.app
import omni.kit.commands
import omni.kit.test
import omni.ui as ui
from omni.ui.tests.test_base import OmniUiTest
from pxr import Kind, Sdf, Gf, Usd


class TestMaterialEditsAndUndo(omni.kit.test.AsyncTestCase):
    def __init__(self, tests=()):
        super().__init__(tests)

    # Before running each test
    async def setUp(self):
        self._usd_context = omni.usd.get_context()
        await self._usd_context.new_stage_async()
        self._stage = self._usd_context.get_stage()
        from omni.kit.test_suite.helpers import arrange_windows

        await arrange_windows(topleft_window="Property", topleft_height=64, topleft_width=800.0)

        renderer = "rtx"
        if renderer not in self._usd_context.get_attached_hydra_engine_names():
            omni.usd.add_hydra_engine(renderer, self._usd_context)        
        
        await self.wait()

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

        await wait_stage_loading()
    
    async def wait(self, updates=3):
        for i in range(updates):
            await omni.kit.app.get_app().next_update_async()

    async def test_l2_material_edits_and_undo(self):
        widget_compare_table = {
            Sdf.ValueTypeNames.Half2.type: (Gf.Vec2h(0.12345, 0.12345), 3),
            Sdf.ValueTypeNames.Float2.type: (Gf.Vec2f(0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Double2.type: (Gf.Vec2d(0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Half3.type: (Gf.Vec3h(0.12345, 0.12345, 0.12345), 3),
            Sdf.ValueTypeNames.Float3.type: (Gf.Vec3f(0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Double3.type: (Gf.Vec3d(0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Half4.type: (Gf.Vec4h(0.12345, 0.12345, 0.12345, 0.12345), 3),
            Sdf.ValueTypeNames.Float4.type: (Gf.Vec4f(0.12345, 0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Double4.type: (Gf.Vec4d(0.12345, 0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Int2.type: Gf.Vec2i(9999, 9999),
            Sdf.ValueTypeNames.Int3.type: Gf.Vec3i(9999, 9999, 9999),
            Sdf.ValueTypeNames.Int4.type: Gf.Vec4i(9999, 9999, 9999, 9999),
            Sdf.ValueTypeNames.UChar.type: 99,
            Sdf.ValueTypeNames.UInt.type: 99,
            Sdf.ValueTypeNames.Int.type: 9999,
            Sdf.ValueTypeNames.Int64.type: 9999,
            Sdf.ValueTypeNames.UInt64.type: 9999,
        }

        from omni.kit.test_suite.helpers import select_prims, wait_stage_loading
        from omni.kit import ui_test

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

        material_path = "/OmniPBR"
        omni.kit.commands.execute("CreateMdlMaterialPrim", mtl_url="OmniPBR.mdl", mtl_name="OmniPBR", mtl_path=material_path)
        
        prim_path = "/OmniPBR/Shader"
        self._usd_context.add_to_pending_creating_mdl_paths(prim_path, True, True)
        await self.wait()

        await select_prims([prim_path])
        # Loading all inputs
        await self.wait()

        all_widgets = ui_test.find_all("Property//Frame/**/.identifier!=''")
        for w in all_widgets:
            id = w.widget.identifier
            if id.startswith("float_slider_"):
                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[13:])
                w.widget.scroll_here_y(0.5)
                await ui_test.human_delay()
                old_value = attr.Get()
                widget_old_value = w.widget.model.get_value()
                new_value = 0.3456

                # Sets value to make sure current layer includes this property
                # so it will not be removed after undo to refresh property 
                # window to resync all widgets.
                await w.input(str(new_value))
                await ui_test.human_delay()
                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[13:])
                self.assertAlmostEqual(attr.Get(), new_value, places=4)
                self.assertAlmostEqual(w.widget.model.get_value(), new_value, places=4)

                omni.kit.undo.undo()
                w.widget.scroll_here_y(0.5)
                self.assertAlmostEqual(attr.Get(), old_value, places=4)
                await self.wait()
                self.assertAlmostEqual(w.widget.model.get_value(), widget_old_value, places=4)
            elif id.startswith("integer_slider_"):
                w.widget.scroll_here_y(0.5)
                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[15:])
                await ui_test.human_delay()
                old_value = attr.Get()
                new_value = old_value + 1

                sdr_metadata = attr.GetMetadata("sdrMetadata")
                if sdr_metadata and sdr_metadata.get("options", None):
                    is_mdl_enum = True
                else:
                    is_mdl_enum = False

                # Sets value to make sure current layer includes this property
                # so it will not be removed after undo to refresh property 
                # window to resync all widgets.
                if is_mdl_enum:
                    attr.Set(old_value)
                    w.widget.model.set_value(new_value)
                else:
                    await w.input(str(new_value))
                await ui_test.human_delay()
                self.assertEqual(attr.Get(), new_value)

                omni.kit.undo.undo()
                await self.wait()
                self.assertEqual(attr.Get(), old_value)
                self.assertEqual(w.widget.model.get_value(), old_value)
            elif id.startswith("drag_per_channel_int"):
                w.widget.scroll_here_y(0.5)
                await ui_test.human_delay()
                sub_widgets = w.find_all("**/IntSlider[*]")
                if sub_widgets == []:
                    sub_widgets = w.find_all("**/IntDrag[*]")
                self.assertNotEqual(sub_widgets, [])
                
                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[17:])
                old_value = attr.Get()

                for child in sub_widgets:
                    child.model.set_value(0)
                    await child.input("9999")
                    await ui_test.human_delay()

                metadata = attr.GetAllMetadata()
                type_name = Sdf.ValueTypeNames.Find(metadata.get(Sdf.PrimSpec.TypeNameKey, "unknown type")).type
                value = widget_compare_table[type_name]
                self.assertEqual(attr.Get(), value)
                
                omni.kit.undo.undo()
                await self.wait()
                self.assertEqual(attr.Get(), old_value)
                self.assertEqual(w.widget.model.get_value(), old_value)
            elif id.startswith("drag_per_channel_"):
                w.widget.scroll_here_y(0.5)
                await ui_test.human_delay()
                sub_widgets = w.find_all("**/FloatSlider[*]")
                if sub_widgets == []:
                    sub_widgets = w.find_all("**/FloatDrag[*]")
                self.assertNotEqual(sub_widgets, [])

                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[17:])
                old_value = attr.Get()

                for child in sub_widgets:
                    await child.input("0.12345")
                    await self.wait()

                metadata = attr.GetAllMetadata()
                type_name = Sdf.ValueTypeNames.Find(metadata.get(Sdf.PrimSpec.TypeNameKey, "unknown type")).type
                value, places = widget_compare_table[type_name]
                self.assertTrue(Gf.IsClose(attr.Get(), value, 1e-3))
                
                # Undo all channels
                for i in range(len(sub_widgets)):
                    omni.kit.undo.undo()
                await self.wait()
                self.assertTrue(Gf.IsClose(attr.Get(), old_value, 1e-04), f"{attr.Get()} != {old_value}")

                for i, child in enumerate(sub_widgets):
                    self.assertAlmostEqual(child.model.get_value()[i], old_value[i], places=places)
            elif id.startswith("bool_"):
                w.widget.scroll_here_y(0.5)
                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[5:])
                await ui_test.human_delay()
                old_value = attr.Get()
                new_value = not old_value 

                # Sets value to make sure current layer includes this property
                # so it will not be removed after undo to refresh property 
                # window to resync all widgets.
                attr.Set(old_value)
                w.widget.model.set_value(new_value)
                await ui_test.human_delay()
                self.assertEqual(attr.Get(), new_value)

                omni.kit.undo.undo()
                await self.wait()
                self.assertEqual(attr.Get(), old_value)
                self.assertEqual(w.widget.model.get_value(), old_value)
            elif id.startswith("sdf_asset_"):
                w.widget.scroll_here_y(0.5)
                attr = self._stage.GetPrimAtPath(prim_path).GetAttribute(id[10:])
                await ui_test.human_delay()
                old_value = attr.Get()
                if isinstance(old_value, Sdf.AssetPath):
                    old_value = old_value.path

                # Skips mdl field
                if old_value.startswith("OmniPBR"):
                    continue

                new_value = "testpath/abc.png"

                # Sets value to make sure current layer includes this property
                # so it will not be removed after undo to refresh property 
                # window to resync all widgets.
                attr.Set(old_value)
                w.widget.model.set_value(new_value)
                await ui_test.human_delay()
                self.assertEqual(attr.Get(), new_value)

                omni.kit.undo.undo()
                await self.wait()
                self.assertEqual(attr.Get(), old_value)
                self.assertEqual(w.widget.model.get_value(), old_value)
