class MonotonicPositiveUnboundedLayer(Module, IMixedEffect):
    def __init__(
        self,
        encodings: Encodings,
        signal_type: str,
        hierarchy_categories: list[str | list[str]] | None = None,
        has_time: bool = False,
        has_signal: bool = False,
        hyperparameters: Hyperparams | None = None,
        non_pos: bool = False,
        non_neg: bool = False,
        non_pos_by_signal: list[bool] | None = None,
        non_neg_by_signal: list[bool] | None = None,
        maximum_strength: float | None = None,
        use_bias: bool | None = None,
        increase_lr: float | None = None,
        name: str | None = None,
    ):
        """Monotonic multiplicative factors affecting sales that also scales ROIs of investments.

        Args:
            hierarchy (pd.DataFrame): The hierarchy that the impact learns on.
            n_instances (int): Number of mixed effect signals. Axis index 2 of effect.
            has_time (bool, optional): Whether the hierarchy is on the time axis. Defaults to False.
            hyperparameters (Hyperparams, optional): Dictionary of hyperparameters for buidling this layer.
            name (str | None, optional): Name of the mixed effect captured. Module parent class sets
                                         to name of class if passes as None.
        """
        super().__init__(hyperparameters=hyperparameters, name=name)
        self.signal_type = signal_type
        self.encodings: Encodings = encodings
        self.has_time = has_time
        self.has_signal = has_signal
        self.hierarchy_categories = hierarchy_categories
        self.non_neg = non_neg
        self.non_pos = non_pos
        self.non_pos_by_signal = non_pos_by_signal
        self.non_neg_by_signal = non_neg_by_signal
        self.maximum_strength = maximum_strength
        self.use_bias = use_bias
        self.increase_lr = increase_lr

    def build(self, input_shapes: InputShapes):
        """Builds the sales_mult hierarchical variable.

        Args:
            input_shapes (InputShapes): The effect and hierarchy shapes.
        """
        self.n_instances = input_shapes.signals[2] if len(input_shapes.signals) > 2 else 1
        if self.has_signal:
            self.n_instances = 1
        shape = [self.n_instances]
        if self.use_bias is None:
            self.use_bias = not (self.has_time or self.has_signal)

        if not self.has_signal:
            signal_enc = self.encodings[self.signal_type]
            if TYPE_CHECKING:
                assert isinstance(signal_enc, Mapping)
            feature_names = tuple(get_lookups(signal_enc))  # pyright: ignore [reportArgumentType]
        else:
            feature_names = ("sales_mult",)

        self.sales_mult = self.hyperparameters.get_submodule(
            "effect_mult",
            module_type=HierchicalEmbedding,
            kwargs=dict(
                encodings=self.encodings,
                columns=self.hierarchy_categories,
                shape=shape,
                use_bias=self.use_bias,
                bias_initializer=0.01,
                increase_lr=self.increase_lr,
                feature_names=feature_names,
            ),
            help="The embedding for the multiplier to apply to each signal before exponentiation.",
        )
        self.use_softplus = self.hyperparameters.get_bool(
            "use_softplus",
            default=True,
            help="Whether to use softplus or exp for the standardization to positive multipliers.",
        )
        self.use_mono = self.hyperparameters.get_bool(
            "use_monotonic", default=False, help="Whether to use a monotonic concave layer to combine signals."
        )
        if self.use_mono:
            self.mono_effect = self.hyperparameters.get_submodule(
                "concave_effect_mult",
                module_type=MonoEffect,
                kwargs=dict(
                    encodings=self.encodings,
                    signal_type=self.signal_type,
                    n_instances=self.n_instances,
                    hierarchy_categories=self.hierarchy_categories,
                    has_signal=self.has_signal,
                ),
                help="Neural Network model that learns the monotonic effect.",
            )

    def __call__(
        self,
        batch: MonotonicPositiveUnboundedInput,
        training: bool = False,
        debug: bool = False,
        skip_metrics: bool = False,
    ) -> MonotonicPositiveUnboundedIntermediaries | DistIntermediaries:
        signals = batch.signals
        if self.use_mono:
            mono_effect_intermediaries = self.mono_effect(
                MonoEffectInput(
                    signals=batch.signals,
                    hierarchy=batch.hierarchy,
                ),
                training=training,
                debug=debug,
                skip_metrics=skip_metrics,
            )
            signals = mono_effect_intermediaries.signals
        # num_gran x num_inst if not has_time and not has_signal
        # num_gran x num_time x num_inst if has_time and not has_signal
        # num_gran x num_inst x 1 if not has_time and has_signal
        # num_gran x num_time x num_inst x 1 if has_time and has_signal
        baseline_sales_effect_raw = self.sales_mult(batch.hierarchy, training=training, skip_metrics=skip_metrics)
        if not self.has_time:
            baseline_sales_effect_raw = tf.expand_dims(baseline_sales_effect_raw, 1)
        if self.has_signal:
            baseline_sales_effect_raw = tf.squeeze(baseline_sales_effect_raw, -1)
        if self.signal_type == "distribution":
            baseline_sales_effect_raw = baseline_sales_effect_raw + tf.constant(3.0, dtype=tf.float32)
        # At this point baseline_sales_effect_raw is always broadcastable to # num_gran x num_time x num_inst
        # learns weightage of each effect signal and applies on it!
        # This is batch x time x n_instances
        # Shifted by -3 to make initialization more sane(before it was very large impacts in initial state).
        softplus_baseline_effect_raw = softplus(baseline_sales_effect_raw - tf.constant(3.0, dtype=tf.float32))
        if self.non_neg:
            baseline_sales_effect_raw = softplus_baseline_effect_raw
        if self.non_pos:
            baseline_sales_effect_raw = -softplus_baseline_effect_raw
        if self.non_pos_by_signal:
            non_pos_by_signal: tf.Tensor = tf.gather(
                tf.constant(self.non_pos_by_signal, dtype=tf.float32, name="non_pos_by_signal"), batch.signal_index
            )
            baseline_sales_effect_raw = baseline_sales_effect_raw * (
                tf.constant(1.0, dtype=tf.float32) - non_pos_by_signal
            ) - (non_pos_by_signal * softplus_baseline_effect_raw)
        if self.non_neg_by_signal:
            non_neg_by_signal = tf.gather(
                tf.constant(self.non_neg_by_signal, dtype=tf.float32, name="non_neg_by_signal"), batch.signal_index
            )
            baseline_sales_effect_raw = (
                baseline_sales_effect_raw * (1 - non_neg_by_signal) + non_neg_by_signal * softplus_baseline_effect_raw
            )
        if self.maximum_strength is not None:
            baseline_sales_effect_raw = (
                tf.math.tanh(baseline_sales_effect_raw / self.maximum_strength) * self.maximum_strength
            )
        baseline_sales_effect = tf.math.multiply(baseline_sales_effect_raw, signals, name="sales_effect")
        baseline_sales_effect = tf.grad_pass_through(lambda x: tf.maximum(-16.0, x, "baseline_sales_effect_clipped"))(
            baseline_sales_effect
        )
        if self.use_softplus:
            # Force 0 to map to 1 after softplus.
            baseline_sales_mult_by_signal = softplus(baseline_sales_effect + np.log(np.e - 1), name="mult_by_signal")
        else:
            baseline_sales_mult_by_signal = tf.math.exp(baseline_sales_effect, name="mult_by_signal")
        # batch x time
        baseline_sales_mult = tf.reduce_prod(baseline_sales_mult_by_signal, 2, name="impact")
        return MonotonicPositiveUnboundedIntermediaries(
            baseline_sales_effect_raw=baseline_sales_effect_raw if debug else None,
            baseline_sales_effect=baseline_sales_effect if debug else None,
            impact_by_signal=baseline_sales_mult_by_signal,
            baseline_sales_mult=baseline_sales_mult if debug else None,
            impact=baseline_sales_mult,
            signal_names=tf.gather(
                tf.convert_to_tensor(get_lookups(self.encodings[self.signal_type])), batch.signal_index
            ),
        )