# 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 carb.settings
import omni.kit.app
import omni.ui as ui
from pxr import Sdf, UsdShade

from ..presets import ColormapPresets
from .index_properties_base_widget import IndexPropertiesBaseWidget

class VolumePropertiesWidget(IndexPropertiesBaseWidget):

    point_cloud_types = [ "PointCloud", "StreamingSettings"]

    def __init__(self):
        self.settings = carb.settings.get_settings()
        super().__init__("IndeX Volume", ["Volume", "Points"] + VolumePropertiesWidget.point_cloud_types)
        self._change_cb = omni.kit.app.SettingChangeSubscription("/rtx/index/compositeEnabled", self._on_change)

    def clean(self):
        self._change_cb = None
        super().clean()

    def _on_change(self, *_):
        self.request_rebuild()

    def set_volume_type(self, combo_model, item):
        prim = self.get_prim()
        volume_type = combo_model.get_current_string()
        if volume_type ==  "default":
            prim.GetAttribute("nvindex:type").Clear()
        else:
            prim.CreateAttribute("nvindex:type", Sdf.ValueTypeNames.Token, True, Sdf.VariabilityUniform).Set(volume_type)
        self.request_rebuild() # update UI to change mode

    def get_compositing(self):
        return self.get_prim().GetAttribute("nvindex:composite").Get() == True

    def set_compositing(self, enable):
        self.get_prim().CreateAttribute("nvindex:composite", Sdf.ValueTypeNames.Bool, custom=True).Set(enable)
        if enable:
            self.set_rtx_skip(True) # don't disable it here
        self.request_rebuild() # update UI

    def get_rtx_skip(self):
        return self.get_prim().GetAttribute("omni:rtx:skip").Get() == True

    def set_rtx_skip(self, skip):
        self.get_prim().CreateAttribute("omni:rtx:skip", Sdf.ValueTypeNames.Bool, custom=True).Set(skip)

    def create_material(self):
        volume_prim = self.get_prim()
        stage = volume_prim.GetStage()
        material_path_base = str(volume_prim.GetPath()) + "/Material"
        material_path = material_path_base

        # Make material path unique
        i = 0
        while stage.GetPrimAtPath(material_path):
            i += 1
            material_path = material_path_base + str(i)

        colormap_path = f"{material_path}/Colormap"
        shader_path = f"{material_path}/VolumeShader"

        colormap = stage.DefinePrim(colormap_path, "Colormap")
        colormap_output = colormap.CreateAttribute("outputs:colormap", Sdf.ValueTypeNames.Token)
        colormap.CreateAttribute("domain", Sdf.ValueTypeNames.Float2).Set((0.0, 1.0))

        material = UsdShade.Material.Define(stage, material_path)
        binding_api = UsdShade.MaterialBindingAPI.Apply(volume_prim)
        binding_api.Bind(material)

        shader = UsdShade.Shader.Define(stage, shader_path)
        shader_output = shader.CreateOutput("volume", Sdf.ValueTypeNames.Token)
        material.CreateVolumeOutput("nvindex").ConnectToSource(shader_output)

        shader.CreateInput("colormap", Sdf.ValueTypeNames.Token).ConnectToSource(colormap_output.GetPath())

        colormapSource = "rgbaPoints"
        colormap.CreateAttribute("colormapSource", Sdf.ValueTypeNames.Token).Set(colormapSource)
        (xPoints, rgbaPoints) = ColormapPresets.generate(ColormapPresets.DEFAULT, colormap_source=colormapSource)
        colormap.CreateAttribute("xPoints", Sdf.ValueTypeNames.FloatArray).Set(xPoints)
        colormap.CreateAttribute("rgbaPoints", Sdf.ValueTypeNames.Float4Array).Set(rgbaPoints)

    def build_properties(self, stage, prim, prim_path):
        with ui.VStack(spacing=5):
            # Volume type
            current_volume_type = prim.GetAttribute("nvindex:type").Get()
            if not current_volume_type and (prim.GetTypeName() == "Points" or prim.GetTypeName() in VolumePropertiesWidget.point_cloud_types):
                current_volume_type = "particle_volume"
                if self._collapsable_frame:
                    self._collapsable_frame.title = "IndeX Points"

            if self.settings.get("/nvindex/compositeRenderingMode") and prim.GetTypeName() not in VolumePropertiesWidget.point_cloud_types:
                if not self.settings.get("/rtx/index/compositeEnabled"):
                    self.build_label(
                            "NVIDIA IndeX Compositing is currently not active. "
                            "You can enable it in 'Render Settings / Common / NVIDIA IndeX Compositing'.", full_width=True)
                    ui.Separator()

                # Enable/disable composite rendering for this prim
                with ui.HStack():
                    self.build_label("Use IndeX compositing")
                    widget = ui.CheckBox(name="nvindex_composite")
                    widget.model.set_value(self.get_compositing())
                    widget.model.add_value_changed_fn(lambda m: self.set_compositing(m.get_value_as_bool()))

                with ui.HStack():
                    self.build_label("Skip in RTX renderer")
                    widget = ui.CheckBox(name="rtx_skip")
                    widget.model.set_value(self.get_rtx_skip())
                    widget.model.add_value_changed_fn(lambda m: self.set_rtx_skip(m.get_value_as_bool()))

                # Skip rest of UI unless compositing is enabled
                if not self.get_compositing():
                    return

                ui.Separator()

            with ui.HStack():
                self.build_label("Volume Type")
                volume_types_list = ["vdb", "sparse_volume"]
                if current_volume_type is None or current_volume_type in volume_types_list:
                    self.volume_type_widget = ui.ComboBox(self.create_combo_list_model(current_volume_type, ["default"] + volume_types_list), name="volume_type")
                    self.volume_type_widget.model.add_item_changed_fn(self.set_volume_type)
                else:
                    self.build_label(current_volume_type)

            render_settings = prim.GetCustomDataByKey(self.USD_RENDERSETTINGS)

            if current_volume_type != "irregular_volume":
                # Common settings (except for irregular volumes)
                self.build_render_setting_float(
                    "Sampling Distance", "samplingDistance",
                    -1.0,
                    render_settings)

                self.build_render_setting_float(
                    "Reference Sampling Distance", "referenceSamplingDistance",
                    -1.0,
                    render_settings)

            if current_volume_type == "sparse_volume":
                # self.build_render_setting_int3(
                #     "Voxel Offsets", "voxelOffsets",
                #     [-1.0, -1.0, -1.0],
                #     render_settings)

                # Filter Mode
                self.build_render_setting_combo_box(
                    "Filter Mode", "filterMode",
                    ["default", "nearest",
                     "trilinear", "trilinear_pre",
                     "tricubic_catmull", "tricubic_catmull_pre",
                     "tricubic_bspline", "tricubic_bspline_pre"],
                    render_settings)

                self.build_render_setting_bool(
                    "LOD Rendering", "lodRenderingEnabled",
                    False,
                    render_settings)

                self.build_render_setting_float(
                    "LOD Pixel Threshold", "lodPixelThreshold",
                    -1.0,
                    render_settings)

                self.build_render_setting_bool(
                    "Preintegration", "preintegratedVolumeRendering",
                    False,
                    render_settings)

            elif current_volume_type == "particle_volume":
                self.build_render_setting_combo_box(
                    "Rendering Mode", "renderingMode",
                    ["default", "rbf", "solid_box", "solid_cone", "solid_cylinder", "solid_sphere", "solid_disc"],
                    render_settings)

                self.build_render_setting_combo_box(
                    "RBF Falloff Kernel", "rbfFalloffKernel",
                    ["default", "constant", "linear", "hermite", "spline"],
                    render_settings)

                # Prefer point width as it is used by USD, but also support IndeX point radius
                if render_settings and "overrideFixedRadius" in render_settings:
                    self.build_render_setting_float(
                        "Override Fixed Radius", "overrideFixedRadius",
                        -1.0,
                        render_settings)
                else:
                    self.build_render_setting_float(
                        "Override Fixed Width", "overrideFixedWidth",
                        -1.0,
                        render_settings)

                if prim.GetTypeName() in VolumePropertiesWidget.point_cloud_types:
                    self.build_render_setting_float(
                        "Spacing Point Width Factor", "spacingPointWidthFactor",
                        -1.0,
                        render_settings)

            elif current_volume_type == "irregular_volume":
                self.build_render_setting_float(
                    "Halo Size", "haloSize",
                    -1.0,
                    render_settings)

                self.build_render_setting_int(
                    "Sampling Mode", "samplingMode",
                    -1,
                    render_settings)

                self.build_render_setting_float(
                    "Sampling Segment Length", "samplingSegmentLength",
                    -1.0,
                    render_settings)

                self.build_render_setting_float(
                    "Sampling Reference Segment Length", "samplingReferenceSegmentLength",
                    -1.0,
                    render_settings)

            else: # vdb (default)
                # Filter Mode
                self.build_render_setting_combo_box(
                    "Filter Mode", "filterMode",
                    ["default", "nearest", "trilinear"],
                    render_settings)

                self.build_render_setting_int(
                    "Subset Voxel Border", "subsetVoxelBorder",
                    -1,
                    render_settings)

        if self.show_advanced_settings:
            if current_volume_type == "irregular_volume":
                self.build_render_setting_int(
                    "Diagnostics Mode", "diagnosticsMode",
                    0,
                    render_settings)

                self.build_render_setting_int(
                    "Diagnostics Flags", "diagnosticsFlags",
                    0,
                    render_settings)

                self.build_render_setting_float(
                    "Wireframe Size", "wireframeSize",
                    0.001,
                    render_settings)

                self.build_render_setting_float(
                    "Wireframe Color Mod Begin", "wireframeColorModBegin",
                    0.0,
                    render_settings)

                self.build_render_setting_float(
                    "Wireframe Color Mod Factor", "wireframeColorModFactor",
                    0.0,
                    render_settings)
            else:
                self.build_render_setting_int(
                    "Debug Visualization", "debugVisualizationOption",
                    0,
                    render_settings)

        # Show timestep only if already defined
        timestep = prim.GetAttribute("nvindex:timestep")
        if timestep:
            with ui.HStack():
                self.build_label("Timestep")
                widget = ui.FloatField(name="timestep")
                widget.model.set_value(timestep.Get() or 0.0)
                widget.model.add_value_changed_fn(lambda m: self.store_value("nvindex:timestep", m.get_value_as_float(), 0.0, Sdf.ValueTypeNames.Double))

        with ui.HStack():
            ui.Button("Create Volume Material", clicked_fn=self.create_material, name="createVolumeMaterial")
