#!/usr/bin/env python3

# Copyright (c) 2022, 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.usd
from omni.kit import ui_test
from pxr import Sdf, UsdShade

# Test custom properties UI for NVIDIA IndeX
class TestIndeXPropertiesUI(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        super().setUp()

        self._usd_context = omni.usd.get_context()

        from omni.hydra.index.tests import USD_DIR
        test_file_path = USD_DIR.joinpath("torus-settings.usda").absolute()
        await self._usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        self._stage = self._usd_context.get_stage()
        self.assertIsNotNone(self._stage)

    def _print_sub_widget(self, parent):
        for w in parent.find_all("**"):
            print("-", w, w.realpath)
            if hasattr(w.widget, 'identifier') and w.widget.identifier:
                print(" id =", w.widget.identifier)
            if hasattr(w.widget, 'name') and w.widget.name:
                print(" name =", w.widget.name)
            if hasattr(w.widget, 'title') and w.widget.title:
                print(" title =", w.widget.title)

    # Test "IndeX Volume" properties of UsdVolume prims
    async def test_volume_properties(self):
        prim = self._stage.GetPrimAtPath("/World/Volume")
        self.assertTrue(prim.IsValid())

        # Select the prim.
        self._usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)
        # Need to wait for additional frames for omni.ui rebuild to take effect
        await ui_test.wait_n_updates(20)

        index_volume = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='IndeX Volume'")
        self.assertTrue(index_volume)

        custom_data = prim.GetCustomData()
        self.assertTrue('nvindex.renderSettings' in custom_data)
        render_settings = custom_data['nvindex.renderSettings']

        filter_mode_widget = index_volume.find("**/ComboBox[0].name=='filterMode'")
        self.assertTrue(filter_mode_widget)
        self.assertEqual(filter_mode_widget.model.get_current_string(), render_settings['filterMode'])

        sampling_distance_widget = index_volume.find("**/FloatField[0].name=='samplingDistance'")
        self.assertTrue(sampling_distance_widget)
        self.assertEqual(sampling_distance_widget.model.get_value_as_float(), render_settings['samplingDistance'])

        # Modify values in UI and check if USD is updated
        filter_mode_widget.model.set_current_index(0) # set to default filter

        # await sampling_distance_widget.input(str(2.345)) # This doesn't seem to work
        sampling_distance_widget.model.set_value(2.345)
        await ui_test.wait_n_updates(2)

        custom_data = prim.GetCustomData()
        render_settings = custom_data['nvindex.renderSettings']
        self.assertTrue('filterMode' not in render_settings) # default value will not be stored
        self.assertEqual(sampling_distance_widget.model.get_value_as_float(), render_settings['samplingDistance'])

        # Now modify the USD and check if UI is updated
        new_filter_mode = "nearest"
        new_sampling_distance = 1.2345
        prim.SetCustomDataByKey("nvindex.renderSettings:filterMode", new_filter_mode)
        prim.SetCustomDataByKey("nvindex.renderSettings:samplingDistance", new_sampling_distance)
        await omni.kit.app.get_app().next_update_async()
        custom_data = prim.GetCustomData()

        # Need to reacquire widgets after rebuild
        filter_mode_widget = index_volume.find("**/ComboBox[0].name=='filterMode'")
        self.assertTrue(filter_mode_widget)
        self.assertEqual(filter_mode_widget.model.get_current_string(), new_filter_mode)

        sampling_distance_widget = index_volume.find("**/FloatField[0].name=='samplingDistance'")
        self.assertTrue(sampling_distance_widget)
        self.assertEqual(sampling_distance_widget.model.get_value_as_float(), new_sampling_distance)

    # Test "IndeX Colormap" properties of (custom) Colormap prims
    async def test_colormap_properties(self):
        prim = self._stage.GetPrimAtPath("/World/Volume/mat/Colormap")
        self.assertTrue(prim.IsValid())

        self._usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)
        await ui_test.wait_n_updates(20)

        index_colormap = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='IndeX Colormap'")
        self.assertTrue(index_colormap)

        # Colormap domain
        domain_group = index_colormap.find("**/drag_per_channel_domain")
        self.assertTrue(domain_group)
        domain_value = []
        for float_drag in domain_group.find_all("**/FloatDrag[0]"):
            domain_value.append(float_drag.widget.model.get_value_as_float())
        self.assertEqual(domain_value, prim.GetAttribute("domain").Get())

        # Colormap domain boundary mode
        domain_boundary_mode_widget = index_colormap.find("**/ComboBox[0].name=='domainBoundaryMode'")
        self.assertTrue(domain_boundary_mode_widget)
        self.assertEqual(domain_boundary_mode_widget.model.get_current_string(), prim.GetAttribute("domainBoundaryMode").Get())

        # Colormap source
        colormap_source_widget = index_colormap.find("**/ComboBox[0].name=='colormapSource'")
        self.assertTrue(colormap_source_widget)
        self.assertEqual(colormap_source_widget.model.get_current_string(), "colormapValues")

        # Set a simple rgbaValues colormap
        prim.CreateAttribute("xPoints", Sdf.ValueTypeNames.FloatArray).Set([0.5, 0.7, 1.0])
        prim.CreateAttribute("rgbaPoints", Sdf.ValueTypeNames.Float4Array).Set([(0, 1, 0, 0.14), (1, 0, 0, 0.25), (1, 1, 0, 0.25)])

        # Switch colormap source
        colormap_source_widget = index_colormap.find("**/ComboBox[0].name=='colormapSource'")
        colormap_source_widget.model.set_current_index(1)
        await ui_test.wait_n_updates(2)

        # Verify values of the first control point in colormap editor
        colormap_point_position_widget = index_colormap.find("**/FloatDrag[*].name=='colormapPointPosition'")
        self.assertTrue(colormap_point_position_widget)
        self.assertEqual(colormap_point_position_widget.model.get_value_as_float(), prim.GetAttribute("xPoints").Get()[0])

        colormap_point_color_widget = index_colormap.find("**/ColorWidget[*].name=='colormapPointColor'")
        self.assertTrue(colormap_point_color_widget)
        color_model = colormap_point_color_widget.model
        sub_models = color_model.get_item_children(None)
        rgba_value = [ color_model.get_item_value_model(s).get_value_as_float() for s in sub_models ]

        colormap_point_opacity_widget = index_colormap.find("**/FloatDrag[*].name=='colormapPointOpacity'")
        self.assertTrue(colormap_point_opacity_widget)
        rgba_value.append(colormap_point_opacity_widget.model.get_value_as_float())

        self.assertEqual(rgba_value, prim.GetAttribute("rgbaPoints").Get()[0])

        # Modify values in UI and check if USD is updated
        colormap_point_position_widget.model.set_value(0.2)
        for s, v in zip(sub_models, [0.1, 0.2, 0.3]):
            color_model.get_item_value_model(s).set_value(v)
        colormap_point_opacity_widget.model.set_value(0.4)
        await ui_test.wait_n_updates(2)

        self.assertEqual(colormap_point_position_widget.model.get_value_as_float(), prim.GetAttribute("xPoints").Get()[0])
        self.assertEqual(colormap_point_opacity_widget.model.get_value_as_float(), prim.GetAttribute("rgbaPoints").Get()[0][3])

        # Load a colormap preset
        load_colormap_preset_widget = index_colormap.find("**/ComboBox[0].name=='loadColormapPreset'")
        self.assertTrue(load_colormap_preset_widget)

        preset_name = "Green Ramp"
        presets = load_colormap_preset_widget.model.get_item_list()
        self.assertTrue(preset_name in presets)
        load_colormap_preset_widget.model.set_current_index(presets.index(preset_name))
        await ui_test.wait_n_updates(2)

        self.assertEqual(prim.GetAttribute("rgbaPoints").Get()[0], [0, 1, 0, 0]) # check that USD was updated

    # Test "IndeX Shader" properties of UsdShader prims
    async def test_shader_properties(self):
        prim = self._stage.GetPrimAtPath("/World/Volume/mat/VolumeShader")
        self.assertTrue(prim.IsValid())

        self._usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)
        await ui_test.wait_n_updates(20)

        index_shader = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='IndeX Shader'")
        self.assertTrue(index_shader)

        # Shader implementation source
        implementation_source_widget = index_shader.find("**/ComboBox[0].name=='info:implementationSource'")
        self.assertTrue(implementation_source_widget)
        self.assertEqual(implementation_source_widget.model.get_current_string(), "id")

        # Provide shader from asset
        implementation_source_widget.model.set_current_index(implementation_source_widget.model.get_item_list().index("sourceAsset"))
        await ui_test.wait_n_updates(2)

        source_asset_widget = index_shader.find("**/sdf_asset_info:xac:sourceAsset.identifier=='sdf_asset_info:xac:sourceAsset'")
        self.assertTrue(source_asset_widget)
        source_asset_widget.model.set_value("./torus-shader.cu")
        await ui_test.wait_n_updates(2)

        shader = UsdShade.Shader(prim)
        self.assertTrue(shader.GetSourceAsset('xac'))

        # Integrate the asset
        integrate_asset_widget = index_shader.find("**/Button[*].name=='integrateAssetContents'")
        self.assertTrue(integrate_asset_widget)
        await integrate_asset_widget.click()
        await ui_test.wait_n_updates(2)
        self.assertTrue(shader.GetSourceCode('xac').startswith("NV_IDX_XAC_VERSION_1_0"))

        # Edit the code
        source_code_widget = index_shader.find("**/StringField[*].name=='info:xac:sourceCode'")
        self.assertTrue(source_code_widget)
        source_code = source_code_widget.model.get_value_as_string() + "// foo bar baz!"
        source_code_widget.model.set_value(source_code)
        await ui_test.wait_n_updates(2)
        self.assertNotEqual(shader.GetSourceCode('xac'), source_code) # not saved yet

        save_compile_widget = index_shader.find("**/Button[*].name=='saveAndCompile'")
        self.assertTrue(save_compile_widget)
        await save_compile_widget.click()
        await ui_test.wait_n_updates(2)
        self.assertEqual(shader.GetSourceCode('xac'), source_code)

        # Load a shader preset
        load_shader_preset_widget = index_shader.find("**/ComboBox[0].name=='loadShaderPreset'")
        self.assertTrue(load_shader_preset_widget)

        preset_name = "Minimal Volume Shader"
        presets = load_shader_preset_widget.model.get_item_list()
        self.assertTrue(preset_name in presets)
        load_shader_preset_widget.model.set_current_index(presets.index(preset_name))
        await ui_test.wait_n_updates(2)
        self.assertTrue(shader.GetSourceCode('xac').startswith("NV_IDX_XAC_VERSION_1_0"))
        self.assertTrue("foo bar baz" not in shader.GetSourceCode('xac'))

    # Test the "Create Volume Material" button in "IndeX Volume" properties of UsdVolume prims
    async def test_create_volume_material(self):
        prim = self._stage.GetPrimAtPath("/World/Volume")
        self.assertTrue(prim.IsValid())

        self._usd_context.get_selection().set_selected_prim_paths([str(prim.GetPath())], True)
        await ui_test.wait_n_updates(20)

        index_volume = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='IndeX Volume'")
        self.assertTrue(index_volume)

        old_material = UsdShade.MaterialBindingAPI(prim).GetDirectBinding().GetMaterial()
        self.assertTrue(old_material)

        # Create new volume material
        create_material_widget = index_volume.find("**/Button[*].name=='createVolumeMaterial'")
        self.assertTrue(create_material_widget)
        await create_material_widget.click()
        await ui_test.wait_n_updates(2)

        # Check that the material was created and bound
        new_material = UsdShade.MaterialBindingAPI(prim).GetDirectBinding().GetMaterial()
        self.assertTrue(new_material)
        self.assertEqual(new_material.GetPath(), str(prim.GetPath()) + "/Material")

        # Restore previous material
        UsdShade.MaterialBindingAPI(prim).Bind(old_material)
