# 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 as ok
import omni.ui as ui
import omni.usd as ou
from .ramp import RampWidget, RampModel


class ColorMap:

    def rebuild_ramp(self):
        self.parent_widget.rebuild()

    def __init__(self,
                 parent_widget,
                 prim,
                 attr_names=["xPoints", "rgbaPoints"],
                 width=400,
                 height=20,
                 height_scalar=50,
                 label="",
                 clamp_values=True):

        self.parent_widget = parent_widget
        self.label = label

        self.selected_model = ui.SimpleIntModel(-1)
        self.selected_model.add_value_changed_fn(self.select_key)

        self.rgba_model = RampModel("rgba", self.selected_model, prim, [attr_names[0], attr_names[1]], [[0], [(1, 1, 1, 1)]], True)
        #self.rgba_model = RampModel(self.selected_model, prim, [attr_names[0], attr_names[1]], [default_positions, default_values], clamp_values)

        self.alpha_model = RampModel("alpha", self.selected_model, prim, [attr_names[0], attr_names[1]], [[0], [(1, 1, 1, 1)]], True)
        #self.alpha_model = RampModel(self.selected_model, prim, [attr_names[0], attr_names[2]], [default_positions, default_scale], clamp_values)
        self.alpha_model.selected.add_value_changed_fn(self.select_key)

        self.add_key_model = ui.SimpleFloatModel(0)
        self.add_key_model.add_value_changed_fn(self.add_key)

        self.remove_key_model = ui.SimpleIntModel(0)
        self.remove_key_model.add_value_changed_fn(self.remove_key)

        self.margin = 2
        self.max_display_intensity = 1

        self.ramp_widget_rgb = None
        self.ramp_widget_alpha = None
        self.intensity_range = None
        self.alpha_range = None

        ok.undo.subscribe_on_change(self.undo_redo_on_change)

        usd_watcher = ou.get_watcher()
        self.usd_subscription = usd_watcher.subscribe_to_change_info_path(
            prim.GetPath(), self.on_usd_changed
        )

        from omni.kit.window.property.templates import LABEL_WIDTH, HORIZONTAL_SPACING

        with self.parent_widget:
            with ui.HStack(spacing=HORIZONTAL_SPACING):
                if label != "":
                    ui.Label(self.label, style={"alignment": ui.Alignment.RIGHT_TOP}, width=LABEL_WIDTH)
                ui.Spacer(width=HORIZONTAL_SPACING)
                with ui.VStack(height=0, spacing=12):
                    with ui.HStack(spacing=HORIZONTAL_SPACING):
                        ui.Label(" ", style={"alignment": ui.Alignment.LEFT})
                    with ui.HStack():
                        ui.Spacer(width=HORIZONTAL_SPACING)
                        self.ramp_widget_alpha = RampWidget(
                            parent_widget,
                            self.alpha_model,
                            self.add_key_model,
                            self.remove_key_model,
                            True,
                            width,
                            height_scalar,
                            scalar_index=RampModel.ALPHA_INDEX,
                            name="alpha")
                    with ui.HStack():
                        ui.Spacer(width=HORIZONTAL_SPACING)
                        self.ramp_widget_rgb = RampWidget(
                            parent_widget,
                            self.rgba_model,
                            self.add_key_model,
                            self.remove_key_model,
                            False,
                            width,
                            height,
                            name="rgb")
                    with ui.HStack(style={"margin": self.margin}):
                        # element to change parameter position
                        ui.Label("position", alignment=ui.Alignment.RIGHT_CENTER)
                        self.position_drag = ui.FloatDrag(width=50, min=0, max=1, name="colormapPointPosition")
                        self.position_changed = self.position_drag.model.subscribe_value_changed_fn(self.set_key_position)
                        self.rgba_model.position = self.position_drag.model
                        self.alpha_model.position = self.position_drag.model

                        # add label for alpha & color selector
                        ui.Label("color", alignment=ui.Alignment.RIGHT_CENTER)
                        
                        # add color selector widget
                        self.value_widget = ui.ColorWidget(0.0, 0.0, 0.0, width=50,name="colormapPointColor")
                        self.sub_value_changed = self.value_widget.model.add_item_changed_fn(self.set_rgb)
                        sub_models = self.value_widget.model.get_item_children()
                        self.rgba_model.value.append(self.value_widget.model.get_item_value_model(sub_models[0]))
                        self.rgba_model.value.append(self.value_widget.model.get_item_value_model(sub_models[1]))
                        self.rgba_model.value.append(self.value_widget.model.get_item_value_model(sub_models[2]))

                        # add opacity float box, sync value in both ramp models
                        ui.Label("opacity", alignment=ui.Alignment.RIGHT_CENTER)
                        self.alpha_drag = ui.FloatDrag(width=60, min=0, max=1, name="colormapPointOpacity")
                        self.alpha_model.value.append(self.alpha_drag.model)
                        self.rgba_model.value.append(self.alpha_drag.model)
                        self.alpha_value_changed = self.alpha_model.value[RampModel.SCALAR_INDEX].subscribe_value_changed_fn(self.set_alpha)

                ui.Spacer(width=HORIZONTAL_SPACING)

        self.update_ui_elements()
        self.selected_model.set_value(0)

# public
    def update_ui_elements(self, *_):
        # print("update_ui_elements")

        # update both ramp models
        new_index_rgba = self.rgba_model.update()
        new_index_intensity = self.alpha_model.update()

        # update values
        if new_index_rgba > 0 or new_index_intensity > 0:
            self.selected_model.set_value(max(new_index_rgba, new_index_intensity))

        # update both ramp widgets
        self.ramp_widget_rgb.update_ui_elements()
        self.ramp_widget_alpha.update_ui_elements()

    def on_shutdown(self):
        ok.undo.unsubscribe_on_change(self.undo_redo_on_change)
        self.value_widget.model.remove_item_changed_fn(self.sub_value_changed)

        # clear widgets & models
        self.position_changed = None
        self.intensity_value_changed = None
        self.alpha_value_changed = None
        self.ramp_widget_rgb = None
        self.rgba_model = None
        self.usd_subscription = None

    def undo_redo_on_change(self, cmds):
        if "IndexSetRampValuesCommand" in cmds:
            self.update_ui_elements()
            self.update_widgets()

# private
    def on_usd_changed(self, *_):
        self.update_ui_elements()
        self.update_widgets()

    def set_key_position(self, *_):
        # print("set_key_position")

        # update order of all value models first
        new_index = self.rgba_model.update_position()
        self.alpha_model.update_position()

        # then update positions
        self.rgba_model.save_and_select_position(new_index)

        if self.ramp_widget_rgb:
            self.ramp_widget_rgb.set_key_position()
        if self.ramp_widget_alpha:
            self.ramp_widget_alpha.set_key_position()

    def set_intensity(self, *_):
        self.alpha_model.update_value(True)
        if self.ramp_widget_alpha:
            self.ramp_widget_alpha.set_key_value()

    def set_alpha(self, *_):
        self.rgba_model.update_value(True, RampModel.ALPHA_INDEX)

    def set_rgb(self, *_):
        self.rgba_model.update_value()

        if self.ramp_widget_rgb:
            self.ramp_widget_rgb.set_key_value()

    def select_key(self, model):
        self.update_widgets()

        index = model.as_int
        if self.ramp_widget_rgb:
            self.ramp_widget_rgb.select_key(index)
        if self.ramp_widget_alpha:
            self.ramp_widget_alpha.select_key(index)

    def add_key(self, model):
        position = model.as_float
        if position < 0:
            return

        self.store_undo_data()

        # add value to all models first
        new_index = self.rgba_model.add_key(position)
        self.alpha_model.add_key(position)

        self.append_to_undo_stack()

        # then update positions
        self.rgba_model.save_and_select_position(new_index)

        self.update_ui_elements()

    def remove_key(self, model):
        if model.as_int < 0:
            return

        self.store_undo_data()

        # remove value from all models first
        new_index = self.rgba_model.remove_selected_key()
        self.alpha_model.remove_selected_key()

        # then update positions
        self.rgba_model.save_and_select_position(new_index)

        self.append_to_undo_stack()

        self.update_ui_elements()

    def update_widgets(self):

        self.value_widget.model.remove_item_changed_fn(self.sub_value_changed)

        # update all values
        self.rgba_model.update_widgets()
        self.rgba_model.update_widgets(True, RampModel.ALPHA_INDEX)
        self.alpha_model.update_widgets(True)

        self.sub_value_changed = self.value_widget.model.add_item_changed_fn(self.set_rgb)

    def store_undo_data(self):
        # print("store undo data")

        self.rgba_model.store_undo_data()
        self.alpha_model.store_undo_data()

    def append_to_undo_stack(self):
        # print("append to undo stack")

        rampValues_rgba = self.rgba_model.append_to_undo_stack()
        rampValues_intensity = self.alpha_model.append_to_undo_stack()

        if rampValues_rgba and rampValues_intensity:
            ok.commands.execute("IndexSetRampValuesCommand", ramps=[rampValues_rgba, rampValues_intensity])
