# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_layers.ipynb.

# %% auto 0
__all__ = ['MonotonicDense']

# %% ../../nbs/00_layers.ipynb 3
from typing import *

from contextlib import contextmanager

import numpy as np
from numpy.typing import ArrayLike, NDArray

import tensorflow as tf
from keras.layers import Dense

# %% ../../nbs/00_layers.ipynb 5
class MonotonicDense(Dense):
  """Monotonic counterpart of the regular Dense Layer of tf.Keras

  Args:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
      indicator_vector: Vector to indicate which of the inputs are monotonically
       increasing or monotonically decreasing or non-monotonic. Has value 1 for
        monotonically increasing and -1 for monotonically decreasing and 0 for 
        non-monotonic variables.
      convexity_indicator: If the value is 0 or 1, then all elements of the 
        activation selector will be 0 or 1, respectevely. If None, epsilon will be
        used to determine the number of 0 and 1 in the activation selector.
      epsilon: Percentage of elements with value 1 in the activation vector if
        convexity_indicator is None, ignored otherwise.
  """

  def __init__(
    self,
    units: int,
    activation: Optional[Union[str, Callable]] = None,
    indicator_vector: Optional[NDArray[np.int_]] = None,
    convexity_indicator: Optional[int] = None,
    epsilon: float = 0.5,
    **kwargs,
  ):
    super(MonotonicDense, self).__init__(units, activation=None, **kwargs)
    self.org_activation = activation

    self.indicator_vector = (
      np.array(indicator_vector, dtype="int32")
      if indicator_vector is not None
      else None
    )
    self.convexity_indicator = convexity_indicator
    self.epsilon = epsilon

  def build(self, input_shape):
    super(MonotonicDense, self).build(input_shape)
    if self.indicator_vector is not None:
      self.indicator_vector = tf.reshape(
        tf.constant(self.indicator_vector),
        (-1, 1),
      )
    if self.convexity_indicator is not None:
      self.convexity_indicator = tf.constant(self.convexity_indicator)

  def get_activation_selector(self, inputs) -> tf.Tensor:
    n_ones = int(round(self.units * self.epsilon))
    n_zeros = self.units - n_ones
    activation_selector = tf.concat(
      [
          tf.ones(shape=(tf.shape(inputs)[0], n_ones), dtype="bool"),
          tf.zeros(shape=(tf.shape(inputs)[0], n_zeros), dtype="bool"),
      ],
      axis=1,
    )
    return activation_selector

  @contextmanager
  def replace_kernel(self):
    """Replace kernel with absolute values of weights based on indicator vector

      This is implemented as context manager to ensure rollback functionality
      after the calculation is finished.
    """
    kernel_org = self.kernel
    
    if self.indicator_vector is not None:

        abs_kernel = tf.abs(self.kernel)
        kernel_1 = tf.where(
            self.indicator_vector == 1,
            abs_kernel,
            self.kernel,
        )
        self.kernel = tf.where(self.indicator_vector == -1, -abs_kernel, kernel_1)
    else:
        self.kernel = tf.abs(self.kernel)
    yield

    self.kernel = kernel_org

  def call(self, inputs):

    # calculate linear forward pass with weights set to either positive or
    # negative values based on the value of indicator_vector
    with self.replace_kernel():
      y = super(MonotonicDense, self).call(inputs)


    if self.org_activation:
      activation = tf.keras.activations.get(self.org_activation)
      if self.convexity_indicator == 1:
        y = activation(y)
      elif self.convexity_indicator == -1:
        y = -activation(-y)
      else:
        activation_selector = self.get_activation_selector(inputs)

        y1 = activation(y)
        y2 = -activation(-y)

        y = tf.where(activation_selector, y1, y2)

    return y

