# 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 asyncio
import functools
from tokenize import Double
import numpy as np
import omni.kit as ok
import omni.ui as ui
import omni.usd as ou
from pxr import Gf

# will be eventually unified with omni.ramp extension

class RampValues:
    def __init__(self, prim_path=None,
                 pos_attr_name=None,
                 val_attr_name=None,
                 old_positions=None,
                 old_values=None,
                 new_positions=None,
                 new_values=None):
        self.prim_path = prim_path
        self.pos_attr_name = pos_attr_name
        self.val_attr_name = val_attr_name
        self.old_pos = old_positions
        self.old_val = old_values
        self.new_pos = new_positions
        self.new_val = new_values


class IndexSetRampValuesCommand(ok.commands.Command):
    def __init__(self, ramps, **kwargs):
        self.ramps = ramps

    def do(self):
        stage = ou.get_context().get_stage()
        if not stage:
            return

        for ramp in self.ramps:
            prim = stage.GetPrimAtPath(ramp.prim_path)
            if not prim:
                return

            if ramp.new_pos and ramp.new_val:
                pos_attr = prim.GetAttribute(ramp.pos_attr_name)
                val_attr = prim.GetAttribute(ramp.val_attr_name)

                if pos_attr and val_attr:
                    pos_attr.Set(ramp.new_pos)
                    val_attr.Set(ramp.new_val)

    def undo(self):
        stage = ou.get_context().get_stage()
        if not stage:
            return

        for ramp in self.ramps:
            prim = stage.GetPrimAtPath(ramp.prim_path)
            if not prim:
                return

            if ramp.old_pos and ramp.old_val:
                pos_attr = prim.GetAttribute(ramp.pos_attr_name)
                val_attr = prim.GetAttribute(ramp.val_attr_name)

                if pos_attr and val_attr:
                    pos_attr.Set(ramp.old_pos)
                    val_attr.Set(ramp.old_val)


ok.commands.register(IndexSetRampValuesCommand)


class RampModel:

    positions_attr_name = ""
    values_attr_name = ""

    default_key_positions = [0]
    default_key_values = [(0, 0, 0, 0)]

    ALPHA_INDEX = 3
    SCALAR_INDEX = 0

    def __init__(self, ramp_type, selected, prim, attr_names, default_keys=None, clamp_values=True):

        self.ramp_type = ramp_type
        self.context = ou.get_context()
        self.prim = prim
        if prim:
            self.prim_path = prim.GetPath()
        else:
            self.prim_path = None

        self.clamp_values = clamp_values

        if type(attr_names) == list:
            self.positions_attr_name = attr_names[0]
            self.values_attr_name = attr_names[1]

        if type(default_keys) == list:
            self.default_key_positions = default_keys[0]
            self.default_key_values = default_keys[1]

        # # # widget models - subscribe to get updates
        # current selected index
        self.selected = selected
        # value on selected index
        self.value = [] * 4
        # position on selected index
        self.position = None
        # show only range up to max_display_value
        self.max_range = None
        
        self.min_value = 0
        self.max_value = 1
        self.max_display_value = 1

        self.key_positions_attr = None
        self.key_values_attr = None

        self.key_positions = []
        self.key_values = []
        
        self.num_keys = 0

        self.range = 1

        self.stage = self.context.get_stage()
        if self.stage:
            if self.prim:
                self.key_positions_attr = self.prim.GetAttribute(self.positions_attr_name)
                self.key_values_attr = self.prim.GetAttribute(self.values_attr_name)

# public
    # call this when the ramp has to sync with the usd stage
    def update(self):
        #print("update")  

        self.stage = self.context.get_stage()
        if not self.stage:
            return -1

        self.prim = self.stage.GetPrimAtPath(self.prim_path)
        if not self.prim:
            return -1

        self.key_positions_attr = self.prim.GetAttribute(self.positions_attr_name)
        self.key_values_attr = self.prim.GetAttribute(self.values_attr_name)

        self.initialize_if_empty_ramp()

        if self.key_positions_attr:
            self.key_positions = list(self.key_positions_attr.Get())
        if self.key_values_attr:
            self.key_values = list(self.key_values_attr.Get())

        self.update_num_keys()
        
        if len(self.key_positions) <= self.selected.as_int:
            return len(self.key_positions) - 1
        if len(self.key_values) <= self.selected.as_int:
            return len(self.key_values) - 1

        return -1
        
    def update_position(self):
        #print("update position") 

        if not self.position:
            return

        index = self.selected.as_int
        position = self.position.as_float

        self.key_positions = list(self.key_positions_attr.Get())
        self.key_values = list(self.key_values_attr.Get())

        # update keys data, order them if user dragged one in front, or behind others
        self.key_positions[index] = position
        new_index = self.order_keys(index, position)
        if new_index != index:
            self.key_values_attr.Set(self.key_values)
            
        return new_index

    def save_and_select_position(self, index):

        self.key_positions_attr.Set(self.key_positions)

        self.select_index(index)

    def update_value(self, is_scalar=False, scalar_index=-1):
        #print("update value: scalar: "+str(is_scalar)+" index: "+str(scalar_index) + " (clamping: "+str(self.clamp_values)+", ramptype: "+self.ramp_type+")")  
        #print("   - self value: "+str(self.value))
        if not self.value:
            return

        # get currently selected key
        index = self.selected.as_int

        if is_scalar:
            if scalar_index < 0:
                value = self.value[RampModel.SCALAR_INDEX].as_float
            else:
                value = self.value[scalar_index].as_float

            # clamp set value
            if self.clamp_values:
                value = max(self.min_value, value)

            if scalar_index < 0:
                self.key_values[index] = value
            else:
                self.key_values[index][scalar_index] = value

            self.update_limits(scalar_index)
        else:            
            if self.clamp_values:
                value_r = max(self.min_value, self.value[0].as_float)
                value_g = max(self.min_value, self.value[1].as_float)
                value_b = max(self.min_value, self.value[2].as_float)

                # get alpha value from RGBA array
                value_alpha = self.key_values[index][RampModel.ALPHA_INDEX]
                self.key_values[index] = Gf.Vec4f(value_r, value_g, value_b, value_alpha)
            else:

                # get alpha value from RGBA array
                value_alpha = self.key_values[index][RampModel.ALPHA_INDEX]
                self.key_values[index] = Gf.Vec4f(self.value[0].as_float, self.value[1].as_float, self.value[2].as_float, value_alpha)

        # update key value entries
        self.key_values_attr.Set(self.key_values)

    def update_limits(self, scalar_index):
        #print("update limits on "+str(self.ramp_type)+" ["+str(self.min_value)+","+str(self.max_value)+"], scalar index: "+str(scalar_index))

        if len(self.key_values) == 0:
            return

        if self.max_range and self.max_range.as_bool:
            self.max_value = 1
            self.range = self.max_display_value - self.min_value
        else:
            if (scalar_index < 0):
                # check ramp type
                if (self.ramp_type == "alpha"):
                    # select max only from alpha entries
                    #self.max_value = max(np.array(self.key_values).T[self.ALPHA_INDEX])

                    # AK: use fixed max range
                    self.max_value = 1.0 
                    
                else:
                    # select max from all values
                    #self.max_value = max(self.key_values)

                    # AK: use fixed range
                    self.max_value = 1.0 
                    
            else:
                #max(np.array(self.key_values).T[scalar_index])
                self.max_value = 1.0 
                
            self.range = self.max_value - self.min_value
        self.range = self.max_value - self.min_value

    def save_position(self, index, position):

        # update widgets
        self.position.set_value(position)

    def save_value(self, value, is_scalar, scalar_index):
        self.key_values = list(self.key_values_attr.Get())
        index = self.selected.as_int

        if is_scalar:            
            # AK: update widgets
            if scalar_index < 0:
                #print(" * ramp.save_value: set scalar index with negative index: "+ str(scalar_index))
                #print("     - retrieving value ("+str(len(self.value))+") at: "+str(RampModel.SCALAR_INDEX))
                #print("*    - value ("+str(value)+") scalar_index: "+str(scalar_index)+ " (is_scalar: "+str(is_scalar)+")")
                #print("     - selected index: "+str(self.selected.as_int))

                # update currently selected value                
                self.value[RampModel.SCALAR_INDEX].set_value(value)
            else:
                #print(" - ramp ("+str(self.ramp_type)+") save_value: set scalar "+str(value)+" index: "+ str(scalar_index) + " (length value: "+str(len(self.value))+")")
                #print("    - sample: "+str(self.value))
                if (self.ramp_type == "alpha"): self.value[0].set_value(value)
                else: self.value[scalar_index].set_value(value)
                
            if self.clamp_values:
                #print(" - ramp: set scalar index + clamp: "+ str(scalar_index))
                if (self.ramp_type == "alpha"): self.value[0].set_value(max(self.min_value, value))
                else: self.value[scalar_index].set_value(max(self.min_value, value))
            else:
                #print(" - ramp: set scalar index, no clamp: "+ str(scalar_index))
                self.value[scalar_index].set_value(value)
        else:   
            #print(" - ramp: save value, no scalar: "+ str(scalar_index))
            # update widgets         
            if self.clamp_values:      
                for i in range(3):      
                    self.value[i].set_value(max(self.min_value, value[i]))
            else:
                for i in range(3): 
                    self.value[i].set_value(value[i])

    # get RGBA value for a given parameter position
    def lookup(self, ct):
        
        # get parameters
        num_key = len(self.key_positions)

        # check default cases
        if (num_key == 1): return self.key_values[0]
        elif (num_key == 0): return [0.0, 0.0, 0.0, 0.0]

        # init parameters
        cur_id = 0

        # iterate to current parameter position
        for i, cur_pos in enumerate(self.key_positions):
            if ct > cur_pos:
                cur_id = i + 1
            else:
                break

        cur_id = min(cur_id, num_key - 1)

        # get values and parameters
        cur_val = self.key_values[cur_id]
        prev_val = self.key_values[max(cur_id-1,0)]
        cur_t = self.key_positions[cur_id]
        prev_t = self.key_positions[max(cur_id-1,0)]
        
        # check parameter positions
        cur_it = 0.0
        if (cur_t != prev_t):
            cur_it = min(max((ct - prev_t) / (cur_t - prev_t), 0.0), 1.0)
        
        if isinstance(self.key_values[0], float):
            # interpolate single instance
            #print(" - lookup interp (ct="+str(ct)+"): "+str(cur_it)+" ("+str(prev_t)+","+str(cur_t)+") with values: prev: "+str(prev_val)+" cur:"+str(cur_val))
            return (1.0 - cur_it) * prev_val + cur_it * cur_val

        elif len(self.key_values[0]) == 4:
            # interpolate rgba values
            itr = (1.0 - cur_it) * prev_val[0] + cur_it * cur_val[0]
            itg = (1.0 - cur_it) * prev_val[1] + cur_it * cur_val[1]
            itb = (1.0 - cur_it) * prev_val[2] + cur_it * cur_val[2]
            ita = (1.0 - cur_it) * prev_val[3] + cur_it * cur_val[3]

            rgba_result = [itr, itg, itb, ita]
            #print(" * lookup RGBA value: "+str(rgba_result))
            return rgba_result

    def add_key(self, position):

        value = 0

        # check if prim attributes are valid
        if self.key_positions_attr.IsValid() and self.key_values_attr.IsValid():
            # lookup parameter position (interpolates RGBA value)
            value = self.lookup(position)
            # print(" - RampModel.add_key: adding value "+str(value)+" at "+str(position)+" to ramp "+str(self.ramp_type))

        self.key_positions = list(self.key_positions_attr.Get())
        self.key_values = list(self.key_values_attr.Get())

        self.key_positions.append(position)
        self.key_values.append(value)

        self.update_num_keys()

        new_index = self.order_keys(self.num_keys - 1, position)

        self.key_values_attr.Set(self.key_values)

        return new_index

    def remove_selected_key(self):
        index = self.selected.as_int
        self.key_positions = list(self.key_positions_attr.Get())

        if len(self.key_positions) > 1:
            self.key_values = list(self.key_values_attr.Get())

            self.key_positions.pop(index)
            self.key_values.pop(index)

            self.update_num_keys()

            self.key_values_attr.Set(self.key_values)

            return index - 1
        else:
            return -1

    def select_index(self, index):
        #print(f"select_index {index}")

        if index < 0:
            return
        
        self.selected.set_value(index)

    def update_widgets(self, is_scalar=False, scalar_index=-1):
        #print("update widgets")

        index = self.selected.as_int
        if index < 0:
            return

        if len(self.key_positions) <= 0 or len(self.key_values) <= 0:
            return

        if self.position:
            self.position.set_value(self.key_positions[index])

        if self.value:
            if is_scalar:
                if scalar_index < 0:
                    #print(" * ramp:update_widget: value length: "+str(len(self.value)) +" SCALAR_INDEX: "+str(RampModel.SCALAR_INDEX))
                    self.value[RampModel.SCALAR_INDEX].set_value(self.key_values[index][3])
                else:
                    num_value = len(self.value)
                    #print("RampModel.update_widget: value length: "+str(num_value) +" index: "+str(scalar_index))
                    #print("  - key value length: "+str(len(self.key_values)) +" index: "+str(index)+" sample: "+str(self.key_values[0]))
                    
                    if (isinstance(self.key_values[0], list)): 
                        self.value[min(scalar_index, num_value-1)].set_value(self.key_values[index][scalar_index])
                    #else: print(" ! RampModel.update_widget: Warning: input is no array type.")
            else:        
                for i in range(3):
                    self.value[i].set_value(self.key_values[index][i])

# private
    def initialize_if_empty_ramp(self):
        if self.key_positions_attr and len(self.key_positions_attr.Get()) == 0 or \
           self.key_values_attr and len(self.key_values_attr.Get()) == 0:
            self.key_positions_attr.Set(self.default_key_positions)
            self.key_values_attr.Set(self.default_key_values)

    def order_keys(self, index, position):
        self.key_positions.sort()

        new_index = self.key_positions.index(position, 0, self.num_keys)
        if new_index != index:
            value = self.key_values.pop(index)
            self.key_values.insert(new_index, value)

        return new_index

    def update_num_keys(self):
        num_positions = len(self.key_positions)
        num_keys = len(self.key_values)
        self.num_keys = min(num_positions, num_keys)

        #print(f"num keys: {self.num_keys}")

    init_key_positions = None
    init_key_values = None

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

        self.init_key_positions = self.key_positions.copy()
        self.init_key_values = self.key_values.copy()

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

        if self.init_key_positions != self.key_positions or \
           self.init_key_values != self.key_values:
            return RampValues(prim_path=self.prim_path,
                                pos_attr_name=self.positions_attr_name,
                                val_attr_name=self.values_attr_name,
                                old_positions=self.init_key_positions,
                                old_values=self.init_key_values,
                                new_positions=self.key_positions,
                                new_values=self.key_values)
        return None


class RampWidget:
    key_handle_scalar_style = {"background_color": 0xFF909090, "border_color": 0xFF000000, "border_width": 1}
    key_handle_scalar_over_range_style = {"background_color": 0xFFDFDFDF, "border_color": 0xFF000000, "border_width": 1}
    key_handle_scalar_selected_style = {"background_color": 0xFF909090, "border_color": 0xFFFFFFFF, "border_width": 1}
    key_handle_scalar_over_range_selected_style = {"background_color": 0xFFDFDFDF, "border_color": 0xFFFFFFFF, "border_width": 1}

    key_handle_outline_style = {"background_color": 0x00000000,
                                "border_color": 0xFF000000,
                                "border_width": 1}
    key_handle_selected_style = {"border_color": 0xFFFFFFFF, "border_width": 2}
    key_handle_style = {"border_color": 0xFF000000, "border_width": 1}

    # attr_names - a list with 2 strings specifiying attribute names (positions, values)
    # default_keys - a list with 2 elements:
    #   key positions - a list of floats
    #   key values - a list of tuples
    def __init__(self,
                 parent_widget,
                 model,
                 model_add,
                 model_remove,
                 is_scalar=False,
                 width=400,
                 height=20,
                 scalar_index=-1,
                 name=""):

        self.parent_widget = parent_widget
        self.model = model
        self.model_add = model_add
        self.model_remove = model_remove
        self.is_scalar = is_scalar
        self.width = width
        self.height = height
        self.scalar_index = scalar_index
        self.name = name

        self.image = None

        self.key_handle_size = 10
        self.key_handle_offset_x = -0.5 * self.key_handle_size

        self.placers_to_indices = {}
        self.indices_to_placers = {}
        self.handles = {}
        
        self.new_placer = None
        self.new_placer_offset_x = None
        self.new_placer_offset_y = None

        self.sub_position_changed = None
        self.sub_value_changed = None

        self.initial_key_index = -1
        self.key_index = -1

        self.width_ratio = 1
        self.height_ratio = 1

        self.async_task_update_keys = None
        self.async_task_add_new_key = None

        # AK: update flag
        #self.do_ramp_update = False

        with ui.ZStack():
            vstack = ui.VStack(height=self.height)
            with vstack:
                self.byte_image_provider = ui.ByteImageProvider()
                self.image = ui.ImageWithProvider(self.byte_image_provider,
                                                  fill_policy=ui.IwpFillPolicy.IWP_STRETCH,
                                                  mouse_pressed_fn=functools.partial(self.on_mouse_button_pressed_2, vstack),
                                                  mouse_moved_fn=functools.partial(self.on_mouse_moved_2, vstack),
                                                  mouse_released_fn=self.on_mouse_button_released_2,
                                                  style={"border_color": 0xFF000000, "border_width": 1})
                self.image.set_computed_content_size_changed_fn(self.update_ui_elements)
            self.zstack_keys = ui.ZStack(height=self.key_handle_size * 2)
                    
        if len(self.indices_to_placers):
            placer = self.indices_to_placers[0]
            self.on_mouse_button_pressed(placer)

        self.update()

# public
    # call this when the ramp has to sync with the usd stage
    def update(self, dt=None):
        #print("update")

        if not self.image:
            return

        self.model.update()
        self.update_ui_elements()

    def on_shutdown(self):
        self.sub_position_changed = None
        self.sub_value_changed = None
        self.image = None

# private
    def basic_clamp(self, num):
        return max(min(num, 1.0), 0.0)

    def update_ui_elements(self):
        #print("update ui elements")

        self.width_ratio = self.image.computed_width / self.width
        self.height_ratio = self.image.computed_height / self.height

        if self.is_scalar:
            self.model.update_limits(self.scalar_index)

        selected_key_index = self.model.selected.as_int
        if selected_key_index >= 0:
            if self.key_index >= 0:
                index = self.key_index
            else:
                index = self.initial_key_index
            if index < 0:
                index = selected_key_index

            self.update_ramp()          
            self.async_task_update_keys = asyncio.ensure_future(self.async_update_keys())

    def set_key_position(self):
        #print("set key positions")

        selected_key_index = self.model.selected.as_int
        if selected_key_index >= 0:
            if self.key_index >= 0:
                index = self.key_index
            else:
                index = self.initial_key_index
            if index < 0:
                index = selected_key_index

            self.update_ramp()
            self.async_task_update_keys = asyncio.ensure_future(self.async_update_keys())

    def set_key_value(self):
        #print("set key values")

        selected_key_index = self.model.selected.as_int
        if selected_key_index >= 0:
            if self.key_index >= 0:
                index = self.key_index
            else:
                index = self.initial_key_index
            if index < 0:
                index = selected_key_index

            self.update_ramp()
            self.async_task_update_keys = asyncio.ensure_future(self.async_update_keys())

    async def async_update_keys(self):
        await ok.app.get_app().next_update_async()

        self.update_keys()

    def on_mouse_button_double_clicked(self, *_):
        #print("delete key")

        selected_key_index = self.model.selected.as_int
        if selected_key_index >= 0:
            self.model_remove.set_value(-1)
            self.model_remove.set_value(selected_key_index)
            
            self.initial_key_index = -1

    def on_mouse_button_pressed_2(self, *arg):
        #print("mouse button pressed - add key")

        self.RELEASED = True

        # add key on the next Kit tick
        if self.async_task_add_new_key:
            self.async_task_add_new_key.cancel()
        self.async_task_add_new_key = asyncio.ensure_future(self.async_add_new_key_on_mouse_button_press(arg))

    def on_mouse_moved_2(self, *arg):
        #print("add key, moved")

        vstack = arg[0]

        # omni.ui bug prevents us from moving the placer right
        # upon its creation, so lets calculate its position x/y
        # that will be used in on_mouse_moved
        cursor_position_x = arg[1]
        self.new_placer_offset_x = cursor_position_x - vstack.screen_position_x

        if self.is_scalar:
            cursor_position_y = arg[2]
            self.new_placer_offset_y = cursor_position_y - vstack.screen_position_y

        self.on_mouse_moved(self.new_placer)

    def update_usd_data(self):

        # update usd values to send to index
        cmap_prim = self.model.prim
        cmap_attr = cmap_prim.GetAttribute("colormapValues")
        cmap_data = cmap_attr.Get()
        num_cmap = len(cmap_data)
        #print(" - properties: colormapValues: "+str(num_cmap)+" ramp type: "+str(self.model.ramp_type))

        key_vals = self.model.key_values
        num_keys = len(key_vals)
        #print(" - key values: "+str(num_keys)+" sample: "+str(key_vals[0]))

        # check if current model has RGBA or alpha data
        if (isinstance(key_vals[0], float)):
            print("RampWidget: updating usd float type values (alpha channel only)")

            # update alpha channel only
            new_data = np.ones((num_cmap, 4))
            for x in range(num_cmap):
                ct = x / num_cmap
                alpha_it = self.model.lookup(ct)
                #print(" - interp alpha "+str(ct)+": "+str(alpha_it))
                new_data[x] = cmap_data[x]
                new_data[x][3] = max(min(alpha_it,1.0),0.0)
            
            cmap_attr.Set(new_data)

        else:
            num_values = len(key_vals[0])
            print("RampWidget: updating usd data with "+str(num_values)+" channel RGBA type")

            # update RGBA color data in usd
            new_data = np.ones((num_cmap, 4))
            for x in range(num_cmap):
                ct = x / num_cmap
                col_it = self.model.lookup(ct)
                new_data[x] = col_it
            
            cmap_attr.Set(new_data)

    def on_mouse_button_released_2(self, *_):

        # write data from control points to usd RGBA field
        # self.update_usd_data() # now handled directly by omni.indexusd

        # if user only clicked on the ramp the self.key_index
        # won't be updated by ::on_mouse_move, use the
        # self.initial_key_index instead
        if self.key_index < 0:
            self.key_index = self.initial_key_index
        if self.key_index < 0:
            self.key_index = self.model.selected.as_int
        if self.key_index < 0:
            self.key_index = 0

        self.new_placer_offset_x = None
        self.new_placer_offset_y = None

    # PRESSED = False

    def on_mouse_button_pressed(self, placer, *_):
        #print("mouse button pressed")

        if self.RELEASED:
            self.initial_placer_position = placer.offset_x.value
            if self.is_scalar:
                self.initial_placer_value = placer.offset_y.value
            
            index = self.placers_to_indices[placer]
            self.model.select_index(index)
            
            self.RELEASED = False

    def select_key(self, selected_key_index):
        #print(f"select key {selected_key_index} in model "+self.model.ramp_type)

        if selected_key_index < 0:
            return

        # establish initial index and buffer index     
        self.initial_key_index = selected_key_index
        self.key_index = -1

        if len(self.handles) == 0:
            return

        if self.model.num_keys != len(self.handles):
            return

        # highlite selected key
        for ii in range(self.model.num_keys):
            if self.initial_key_index == ii:
                if self.is_scalar:
                    key_style = self.key_handle_scalar_selected_style
                    if self.model.max_range and self.model.max_range.as_bool:
                        if self.scalar_index < 0:
                            scalar_value = self.model.key_values[ii]
                        else:
                            scalar_value = self.model.key_values[ii][self.scalar_index]

                        if scalar_value > self.model.max_display_value:
                            key_style = self.key_handle_scalar_over_range_selected_style
                    self.handles[ii].set_style(key_style)
                else:
                    self.handles[ii][2].set_style(self.key_handle_selected_style)
                    self.handles[ii][3].set_style(self.key_handle_selected_style)
            else:
                if self.is_scalar:
                    key_style = self.key_handle_scalar_style
                    if self.model.max_range and self.model.max_range.as_bool:
                        if self.scalar_index < 0:
                            scalar_value = self.model.key_values[ii]
                        else:
                            scalar_value = self.model.key_values[ii][self.scalar_index]

                        if scalar_value > self.model.max_display_value:
                            key_style = self.key_handle_scalar_over_range_style
                    self.handles[ii].set_style(key_style)
                else:
                    self.handles[ii][2].set_style(self.key_handle_style)
                    self.handles[ii][3].set_style(self.key_handle_style)

    async def async_add_new_key_on_mouse_button_press(self, *arg):
        await ok.app.get_app().next_update_async()

        if self.RELEASED:
            #print("async add key")

            vstack = arg[0][0]
            cursor_position_x = arg[0][1]

            position = (cursor_position_x - vstack.screen_position_x) / self.image.computed_width
            self.model_add.set_value(-1)
            self.model_add.set_value(position)

            self.RELEASED = False

    async def async_update_keys_on_mouse_button_release(self, index):
        await ok.app.get_app().next_update_async()
        #print("asymc mouse button released")

        self.update_keys()

        self.RELEASED = True

    RELEASED = False

    def on_mouse_button_released(self, *_):
        #print("mouse button released")

        self.RELEASED = True

        # remap ui widgets if user reordered keys
        if self.initial_key_index >= 0 and self.key_index >= 0 and self.initial_key_index != self.key_index:
            placer1 = self.indices_to_placers[self.initial_key_index]
            placer2 = self.indices_to_placers[self.key_index]
            self.placers_to_indices[placer1] = self.key_index
            self.placers_to_indices[placer2] = self.initial_key_index
            self.indices_to_placers[self.initial_key_index] = placer2
            self.indices_to_placers[self.key_index] = placer1

            handles1 = self.handles[self.initial_key_index]
            handles2 = self.handles[self.key_index]
            self.handles[self.initial_key_index] = handles2
            self.handles[self.key_index] = handles1

        if self.is_scalar:
            self.model.update_limits(self.scalar_index)

        # update keys on the next tick
        if self.async_task_update_keys:
            self.async_task_update_keys.cancel()
        
        if self.key_index >= 0:
            index = self.key_index
        else:
            index = self.initial_key_index
        if index < 0:
            index = self.model.selected.as_int

        self.async_task_update_keys = asyncio.ensure_future(self.async_update_keys_on_mouse_button_release(index))

        # restore initial conditions
        self.initial_key_index = -1
        self.key_index = -1
        self.new_placer = None

        self.RELEASED = False

    def on_mouse_moved(self, *arg):
        #print("move key")

        placer = arg[0]
        y = 0

        if self.new_placer_offset_x is None:
            x = placer.offset_x.value
        else:
            x = self.new_placer_offset_x

        if self.is_scalar:
            # Workaround to omni.ui bug where new placers
            # cannot be moved right after their creation.
            # If moving existing key use it's placer position,
            # otherwise use the values provided by on_mouse_moved_2.

            if self.new_placer_offset_y is None:
                y = placer.offset_y.value
            else:
                y = self.new_placer_offset_y

            key_scalar_value = (self.height - y) * (self.model.range / self.height)
            
            if key_scalar_value < self.model.min_value:
                key_scalar_value = self.model.min_value
            # elif key_scalar_value > self.model.max_value:
            #     key_scalar_value = self.model.max_value

            self.model.save_value(key_scalar_value, self.is_scalar, self.scalar_index)

        width = self.width * self.width_ratio
        key_position = round(max(0, min(1, x / width)), 4)

        if self.key_index < 0:
            self.key_index = self.initial_key_index

        self.model.save_position(self.key_index, key_position)

    def update_ramp(self):
        #print("update ramp")

        self.model.initialize_if_empty_ramp()

        num_pixels_x = self.width
        num_pixels_y = self.height if self.is_scalar else 1

        # initialize pixel values
        pixels = [1] * num_pixels_x * num_pixels_y * 4

        if self.model.key_positions_attr.IsValid() and self.model.key_values_attr.IsValid():
            if self.is_scalar:
                if self.scalar_index == RampModel.ALPHA_INDEX:
                    # pixels = self.ramp.get_alpha_array_as_rgba8(self.width, self.height,
                    #                                             self.model.min_value, self.model.max_value,
                    #                                             self.model.key_positions_attr.GetPath().pathString,
                    #                                             self.model.key_values_attr.GetPath().pathString)
                    #self.model.key_positions
                    #print("creating alpha pixel map")

                    # create pixel pattern
                    if (True):
                        # AK: create pixel pattern
                        for x in range(num_pixels_x):
                            # interpolate alpha value
                            xt = float(x / num_pixels_x)
                            cur_rgba = self.model.lookup(xt)
                            cur_alpha = cur_rgba[RampModel.ALPHA_INDEX]
                            max_y = (1.0 - cur_alpha) * num_pixels_y
                            # fill column
                            for y in range(num_pixels_y):
                                index = y * num_pixels_x * 4 + x * 4
                                if (max_y < y):
                                    pixels[index + 0] = 0.5
                                    pixels[index + 1] = 0.5
                                    pixels[index + 2] = 0.5
                                    pixels[index + 3] = cur_alpha
                                else:
                                    pixels[index + 0] = 0.0
                                    pixels[index + 1] = 0.0
                                    pixels[index + 2] = 0.0
                                    pixels[index + 3] = 0.0

                elif self.scalar_index < 0:
                    # pixels = self.ramp.get_float_array_as_rgba8(self.width, self.height,
                    #                                             self.model.min_value, self.model.range - self.model.min_value,
                    #                                             self.model.key_positions_attr.GetPath().pathString,
                    #                                             self.model.key_values_attr.GetPath().pathString)
                    #print("creating float pixel map (sample: "+str(self.model.key_values[0])+")")
                    if (True):
                        # AK: create pixel pattern
                        for x in range(num_pixels_x):
                            # interpolate alpha value
                            xt = float(x / num_pixels_x)
                            cur_rgba = self.model.lookup(xt)
                            cur_alpha = cur_rgba[RampModel.ALPHA_INDEX]
                            max_y = (1.0 - cur_alpha) * num_pixels_y
                            # fill column
                            for y in range(num_pixels_y):
                                index = y * num_pixels_x * 4 + x * 4
                                if (max_y < y):
                                    pixels[index + 0] = 0.5
                                    pixels[index + 1] = 0.5
                                    pixels[index + 2] = 0.5
                                    pixels[index + 3] = cur_alpha
                                else:
                                    pixels[index + 0] = 0.0
                                    pixels[index + 1] = 0.0
                                    pixels[index + 2] = 0.0
                                    pixels[index + 3] = 0.0
                        
                        # set flag
                        #self.do_ramp_update = False
            else:
                positions = [x / (num_pixels_x - 1) for x in range(num_pixels_x)]
                #print("creating generic pixel map from values (#positions: "+str(len(positions))+", sample: "+str(positions[1])+")")
                #print(" - num-x: "+str(num_pixels_x)+" num-y: "+str(num_pixels_y))
                #print(" - length values: "+str(len(self.model.key_values)))
                #print(" - positions: "+str((self.model.key_positions)))
                #print(" - values: "+str((self.model.key_values)))
                
                num_values = len(self.model.key_values)
                num_positions = len(self.model.key_positions) #FIXME: find out why these are not always the same?!??
                pidx = 0
                # values = self.ramp.get_float4_array(positions,
                #                                     self.model.key_positions_attr.GetPath().pathString,
                #                                     self.model.key_values_attr.GetPath().pathString, "")

                if True:
                    for x in range(num_pixels_x):
                        xt = x / num_pixels_x

                        key_pos = self.model.key_positions[pidx]
                        key_pos_prev = self.model.key_positions[max(pidx-1,0)]
                        if (key_pos != key_pos_prev):
                            key_it = min(max((xt - key_pos_prev) / (key_pos - key_pos_prev),0.0),1.0)
                        else: key_it = 0.0

                        key_value_next = self.model.key_values[pidx]
                        key_value_prev = self.model.key_values[max(pidx-1,0)]
                        
                        itr = (1.0 - key_it) * key_value_prev[0] + (key_it) * key_value_next[0]
                        itg = (1.0 - key_it) * key_value_prev[1] + (key_it) * key_value_next[1]
                        itb = (1.0 - key_it) * key_value_prev[2] + (key_it) * key_value_next[2]
                        ita = (1.0 - key_it) * key_value_prev[3] + (key_it) * key_value_next[3]
                        #print(" ramp: it: "+str(key_it))

                        if key_pos < xt and pidx + 1 < min(num_values, num_positions):
                            pidx += 1
                        
                        for y in range(num_pixels_y):
                            
                            # pixels[index + 0] = values[x][0]
                            # pixels[index + 1] = values[x][1]
                            # pixels[index + 2] = values[x][2]
                            # pixels[index + 3] = values[x][3]
                            
                            index = y * num_pixels_x * 4 + x * 4
                            pixels[index + 0] = itr
                            pixels[index + 1] = itg
                            pixels[index + 2] = itb
                            pixels[index + 3] = 1.0

        self.byte_image_provider.set_bytes_data(pixels, [num_pixels_x, num_pixels_y])

    def update_keys(self):
        #print("update keys")

        # currently omni.ui API cannot delete widgets - hide them instead
        for placer in self.indices_to_placers.values():
            placer.visible = False

        for handle in self.handles.values():
            if self.is_scalar:
                handle.visible = False
            else:
                handle[0].visible = False
                handle[1].visible = False
                handle[2].visible = False
                handle[3].visible = False

        self.placers_to_indices = {}
        self.indices_to_placers = {}
        self.handles = {}

        self.model.update_num_keys()
        
        with self.zstack_keys:
            for i in range(self.model.num_keys):
                x = self.model.key_positions[i] * self.width
                value = None
  
                if self.is_scalar:
                    if self.model.range:
                        if self.scalar_index < 0:
                            scalar_value = self.model.key_values[i]
                        else:
                            scalar_value = self.model.key_values[i][self.scalar_index]

                        if self.model.max_range and self.model.max_range.as_bool:
                            # the key will be clamped to max display value
                            scalar_value = min(scalar_value, self.model.max_display_value)

                        y = self.height - abs((scalar_value - self.model.min_value) / self.model.range * self.height)
                    else:
                        y = self.height
                else:
                    y = 0
                    value = self.model.key_values[i]

                self.create_key_handle(i, x, y, value)

        self.select_key(self.model.selected.as_int)

    def create_key_handle(self, index, x, y, value):
        
        offset = self.key_handle_offset_x
        with ui.Placer(offset_x=offset, offset_y=offset):
            placer = ui.Placer(width=self.key_handle_size,
                               height=self.key_handle_size,
                               offset_x=x * self.width_ratio,
                               offset_y=y * self.height_ratio,
                               draggable=True,
                               drag_axis=ui.Axis.XY if self.is_scalar else ui.Axis.X,
                               stable_size=True)
            placer.set_offset_x_changed_fn(functools.partial(self.on_mouse_moved, placer))
            if self.is_scalar:
                placer.set_offset_y_changed_fn(functools.partial(self.on_mouse_moved, placer))

            with placer:
                with ui.ZStack(mouse_pressed_fn=functools.partial(self.on_mouse_button_pressed, placer),
                               mouse_released_fn=self.on_mouse_button_released):
                    if self.is_scalar:
                        circle = ui.Circle(radius=self.key_handle_size * 0.5,
                                           mouse_double_clicked_fn=self.on_mouse_button_double_clicked, 
                                           style=self.key_handle_scalar_style)

                        self.handles[index] = circle
                    else:
                        with ui.Placer():
                            border_triangle = ui.Triangle(
                                width=self.key_handle_size,
                                height=self.key_handle_size * 0.5,
                                alignment=ui.Alignment.CENTER_TOP,
                                style=self.key_handle_outline_style
                            )
                        with ui.Placer(offset_y=self.key_handle_size * 0.5 + 0.5):
                            border_rectangle = ui.Rectangle(width=self.key_handle_size,
                                                            height=self.key_handle_size - 2,
                                                            style=self.key_handle_outline_style)

                        r = int(self.basic_clamp(value[0]) * 255)
                        g = int(self.basic_clamp(value[1]) * 255) << 8
                        b = int(self.basic_clamp(value[2]) * 255) << 16
                        a = 255 << 24

                        with ui.Placer(offset_x=1, offset_y=1):
                            triangle = ui.Triangle(
                                width=self.key_handle_size - 2,
                                height=self.key_handle_size * 0.5 - 1,
                                mouse_double_clicked_fn=self.on_mouse_button_double_clicked,
                                alignment=ui.Alignment.CENTER_TOP,
                                style={"background_color": (r + g + b + a)}
                            )
                        with ui.Placer(offset_x=1, offset_y=self.key_handle_size * 0.5 - 0.5):
                            rectangle = ui.Rectangle(
                                width=self.key_handle_size - 1.75,
                                height=self.key_handle_size - 2,
                                mouse_double_clicked_fn=self.on_mouse_button_double_clicked,
                                style={"background_color": (r + g + b + a)}
                            )

                        self.handles[index] = (triangle, rectangle, border_triangle, border_rectangle)

            self.placers_to_indices[placer] = index
            self.indices_to_placers[index] = placer

        return placer

