# 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
import omni.ui as ui
import omni.usd as ou

from pxr import Sdf
from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidgetBuilder

from ..presets import ColormapPresets
from .index_properties_base_widget import IndexPropertiesBaseWidget

from .colormap import ColorMap
from omni.kit.property.usd.custom_layout_helper import CustomLayoutFrame

class ColormapPropertiesWidget(IndexPropertiesBaseWidget):

    USD_RENDERSETTINGS = 'nvindex.renderSettings'

    def __init__(self):
        super().__init__("IndeX Colormap", "Colormap")
        self.colormap_source_combo_box = None

    def switch_colormap_source(self, value):
        key = "colormapSource"
        self.request_rebuild() # update UI to change mode
        self.store_value(key, value, "colormapValues")

    def load_colormap_preset(self, prim, preset_name, colormapSource):
        if prim and preset_name:
            if colormapSource == "rgbaPoints":
                (xPoints, rgbaPoints) = ColormapPresets.generate(preset_name, colormap_source=colormapSource)

                xPoints_attr = prim.CreateAttribute("xPoints", Sdf.ValueTypeNames.FloatArray)
                xPoints_attr.Set(xPoints)

                rgbaPoints_attr = prim.CreateAttribute("rgbaPoints", Sdf.ValueTypeNames.Float4Array)
                rgbaPoints_attr.Set(rgbaPoints)
            else:
                colormap_data = ColormapPresets.generate(preset_name)
                colormap_attr = prim.CreateAttribute("colormapValues", Sdf.ValueTypeNames.Color4fArray)
                colormap_attr.Set(colormap_data)

    def build_properties(self, stage, prim, prim_path):
        with ui.VStack(spacing=5):
            stage = ou.get_context().get_stage()
            node_prim_path = self._payload[-1]

            #TODO: Make this explicit to handle missing attributes
            self.domain_model = UsdPropertiesWidgetBuilder._vec2_per_channel_builder(
                stage,
                "domain",
                Sdf.ValueTypeNames.Float2,
                { Sdf.PropertySpec.DisplayNameKey: 'Domain' }, # metadata
                [ node_prim_path ]
            )

            with ui.HStack():
                key = "domainBoundaryMode"
                with ui.HStack():
                    self.build_label("Domain Boundary Mode")
                    value = prim.GetAttribute(key).Get() or ""
                    self.domain_boundary_mode_combo_box = ui.ComboBox(self.create_combo_list_model(value, ['clampToEdge', 'clampToTransparent']), name=key)
                    self.domain_boundary_mode_combo_box.model.add_item_changed_fn(lambda m, i: self.store_value("domainBoundaryMode", m.get_current_string(), 'clampToEdge', Sdf.ValueTypeNames.Token))

            with ui.HStack():
                key = "colormapSource"
                with ui.HStack():
                    self.build_label("Colormap Source")
                    value = prim.GetAttribute(key).Get() or ""
                    self.colormap_source_combo_box = ui.ComboBox(self.create_combo_list_model(value, ['colormapValues', 'rgbaPoints']), name=key)
                    self.colormap_source_combo_box.model.add_item_changed_fn(lambda m, i: self.switch_colormap_source(m.get_current_string()))

            colormapSource = prim.GetAttribute("colormapSource").Get() or "colormapValues"
            if colormapSource == "rgbaPoints":
                # adding custom colormap editor
                frame = CustomLayoutFrame(hide_extra=True)
                with frame:
                    colormap_frame = ui.Frame(height=0)
                    with colormap_frame:
                        ColorMap(frame, prim, ["xPoints", "rgbaPoints"])

            with ui.HStack():
                self.build_label("Load Colormap Preset")
                preset_combo_box = ui.ComboBox(self.create_combo_list_model('', [''] + ColormapPresets.get_labels()), name='loadColormapPreset')
                preset_combo_box.model.add_item_changed_fn(lambda model, i: self.load_colormap_preset(prim, model.get_current_string(), colormapSource))
