def custom_metric(func: CustomMetricFunc) -> CustomMetricFunc:
    """Denotes a custom metric function."""
    # TODO (legendof-selda) check if check metric is following the same Protocol
    global CUSTOM_METRICS
    CUSTOM_METRICS[func.__name__] = func
    return func

@custom_metric
def roi_variation(outputs: EconomicIntermediaries, mask: TensorLike | None = None) -> tf.Tensor:  # noqa: U100
    """Calculate the mean absolute log of overall rois as metric which gets amplified whenever
       some rois are extremely high or extremely low for any granularity. Requires time axis at 0.

    Args:
        outputs (EconomicIntermediaries): Outputs of the Economic model.
        mask (TensorLike, optional): Mask is not used. Created to ensure signature of custom_metric function.

    Returns:
        tf.Tensor: The mean absolute log values of overall rois for each granularity
                    with a resulting shape of (batch,).
    """
    total_impacts = (
        outputs.impacts.roicurve.impact_by_signal_instant * outputs.impacts.roicurve.betagamma.decayed_impact_mult
    )
    spends = outputs.inputs.vehicle_spends
    # Time axis reduced
    spends = tf.reduce_sum(spends, axis=1)
    overall_roi = tf.reduce_sum(total_impacts, axis=1) / (spends + EPSILON)
    non_zero_spends_mask = spends > 0
    # Vehicle axis reduced
    return weighted_mean(np.square(np.log(overall_roi + EPSILON)), axis=1, mask=non_zero_spends_mask)

def weighted_custom_metric(func: CustomMetricFunc) -> WeightedCustomMetricFunc:
    """Wraps a custom metric function so that it can calculate weighted average on the custom metric."""

    @wraps(func)
    def wrapped(
        outputs: EconomicIntermediaries,
        mask: TensorLike | None,
        weights: TensorLike | None = None,
        axis: int = 0,
        keepdims: bool = False,
    ) -> tf.Tensor:
        f"""Calculate the weighted {func.__name__} metric.

        Args:
            outputs (EconomicIntermediaries): Outputs of the Economic model.
            mask (TensorLike): Mask to be applied on the time series. Defaults to None.
            weights (TensorLike | None): An array of weights associated with the values in `custom_metric_result`.
                Each value in `custom_metric_result` contributes to the average according to its associated weight.
                The weights array can either be 1-D (in which case its length must be the size of a along the given
                axis) or of the same shape as a.
                If weights=None, then we don't perform weighted average and return the `custom_metric_result`.
                Defaults to None.
            axis (int, optional): Axis along which metrics are computed. Defaults to 0.
            keepdims (bool, optional): Whether to keep the dimensions we take the mean over. Defaults to False.

        Returns:
            tf.Tensor: Weighted {func.__name__} calculated along `axis`.
        """
        custom_metric_result = func(outputs, mask=mask)
        if weights is None:
            return custom_metric_result
        else:
            return weighted_mean(custom_metric_result, weights, axis=axis, keepdims=keepdims)

    return wrapped