@metric_with_variants
def log_cosh(
    y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None = None, axis: int = 0, name: str = "Mae"
) -> tf.Tensor:
    """Calculate the log of the cosh of the error with a temperature factor to control the
    transition between the quadratic and linear regions.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The mean log cosh of the error calculated along `axis`.
    """
    TEMPERATURE = 1e4
    with tf.name_scope(name) as scope:
        abs_residual = tf.abs(residuals(y_true, y_pred, mask)) / TEMPERATURE
        return TEMPERATURE * weighted_mean(
            abs_residual + tf.math.softplus(-2.0 * abs_residual) - tf.math.log(2.0), mask, axis=axis, name=scope
        )

@metric_with_variants
def mae(
    y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None = None, axis: int = 0, name: str = "Mae"
) -> tf.Tensor:
    """Calculate the mean absolute error which is the mean of the absolute value of the residuals.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The mean absolute error calculated along `axis`.
    """
    with tf.name_scope(name) as scope:
        return weighted_mean(tf.abs(residuals(y_true, y_pred, mask)), mask, axis=axis, name=scope)

@metric_with_variants
def mape(
    y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None = None, axis: int = 0, name: str = "Mape"
) -> tf.Tensor:
    """Calculate the mean absolute percentage error. This is the mean of the residuals
       divided by the true values.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The mean absolute percentage error calculated along `axis`.
    """
    with tf.name_scope(name) as scope:
        if mask is None:
            mask = tf.ones_like(y_true)
        return weighted_mean(tf.abs(residuals(y_true, y_pred, mask) / (y_true + EPSILON)), mask, axis=axis, name=scope)

@metric
def mase(
    y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None = None, axis: int = 0, name: str = "Mase"
) -> tf.Tensor:
    """
    Computes the Mean Absolute Scaled Error (MASE).

    MASE is a metric used to evaluate forecast accuracy. It measures the performance
    of a forecasting method against a naive baseline, considering the scale of the time series.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: Calculated Mean Absolute Scaled Error.
    """
    with tf.name_scope(name) as scope:
        mask = tf.convert_to_tensor(mask, dtype_hint=tf.bool, name="mask")
        numeric_mask = tf.cast(mask, dtype=y_true.dtype, name="numeric_mask")
        naive_mae = tf.abs(tf.experimental.numpy.diff(y_true * numeric_mask, axis=axis))
        weighted_naive_mae = tf.math.divide_no_nan(
            tf.reduce_sum(naive_mae, axis=axis),
            tf.reduce_sum(numeric_mask, axis=axis),
        )
        return mae(y_true, y_pred, mask, axis, scope) / (weighted_naive_mae + EPSILON)

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

@metric_with_variants
def mse(y_true: TensorLike, y_pred: TensorLike, mask: TensorLike, axis: int = 0, name: str = "MSE") -> tf.Tensor:
    """Calculate the mean squared error which is the mean of the square of the residuals.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The mean squared error calculated along `axis`.
    """
    with tf.name_scope(name) as scope:
        return weighted_mean(tf.square(residuals(y_true, y_pred, mask)), mask, axis=axis, name=scope)

def residuals(y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None) -> tf.Tensor:
    """Calculate the residuals which are the difference between the true and forecased values.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike | None): The values which should be masked out of the calculation.

    Returns:
        tf.Tensor: The masked absolute difference between y_true and y_pred.
    """
    res = y_true - y_pred
    return res if mask is None else res * tf.cast(mask, dtype=res.dtype)

@metric_with_variants
def rmse(y_true: TensorLike, y_pred: TensorLike, mask: TensorLike, axis: int = 0, name: str = "RMSE") -> tf.Tensor:
    """Calculate the root mean squared error which is the square root of the mean squared error.
       We take the square root so the result has the same units as the input.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The root mean squared error calculated along `axis`.
    """
    with tf.name_scope(name) as scope:
        return tf.sqrt(mse(y_true, y_pred, mask, axis=axis, name=scope))

@metric_with_variants
def smape(
    y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None = None, axis: int = 0, name: str = "SMape"
) -> tf.Tensor:
    """Calculate the symmetric mean absolute percentage error. This calculates the percentage
       errors as twice the residual divided by the sum of forecasted and true values.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The symmetric mean absolute percentage error calculated along `axis`.
    """
    with tf.name_scope(name) as scope:
        return weighted_mean(
            2 * tf.abs(residuals(y_true, y_pred, mask) / (tf.abs(y_true) + tf.abs(y_pred) + EPSILON)),
            mask,
            axis=axis,
            name=scope,
        )

def weighted_mean(
    data: TensorLike,
    mask: TensorLike | None,
    axis: list[int] | tuple[int, ...] | int | None = None,
    keepdims: bool = False,
    name: str = "WeightedMean",
):
    """Calculate the mean of the input array weighted by the mask along axis.
        Requires time axis at 0.
    Args:
        data (TensorLike): The array to take the mean over.
        mask (TensorLike | None): The weights (or mask if it is all 0's and 1's) to use when taking the mean.
                None would simply take mean over the given axis.
        axis (int | Sequence[int] | None, optional): The axis to take the mean over. Defaults to None.
        keepdims (bool, optional): Whether to keep the dimensions we take the mean over. Defaults to False.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The weighted mean of data.
    """
    with tf.name_scope(name) as scope:
        if mask is None:
            return tf.reduce_mean(data, axis=axis, keepdims=keepdims)

        data = tf.convert_to_tensor(data, dtype_hint=data.dtype, name="data")
        mask = tf.convert_to_tensor(mask, dtype_hint=tf.bool, name="mask")
        numeric_mask = tf.cast(mask, dtype=data.dtype, name="numeric_mask")
        return tf.math.divide_no_nan(
            tf.math.reduce_sum(data * numeric_mask, axis=axis, keepdims=keepdims),
            tf.math.reduce_sum(numeric_mask, axis=axis, keepdims=keepdims),
            name=scope,
        )

def weighted_metric(func: MetricFunc) -> WeightedMetricFunc:
    """Wraps a metric function so that it can calculate weighted average on the metric."""

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

        Args:
            y_true (TensorLike): The true values of the time series.
            y_pred (TensorLike): The forecasted values of the time series.
            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 `metric_result`.
                Each value in `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 `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.
            name (str, optional): Python str name prefixed to ops created by this function.

        Returns:
            tf.Tensor: Weighted {func.__name__} calculated along `axis`.
        """
        with tf.name_scope(name) as scope:
            metric_result = func(y_true, y_pred, mask, axis, scope)
            if weights is None:
                return metric_result
            else:
                return weighted_mean(
                    metric_result, weights, axis=axis - 1 if axis > 0 else -1, keepdims=keepdims, name=scope
                )

    return wrapped

@metric_with_variants
def wmape(
    y_true: TensorLike, y_pred: TensorLike, mask: TensorLike | None = None, axis: int = 0, name: str = "WMape"
) -> tf.Tensor:
    """Calculate the weighted mean absolute percentage error. This is the mean of the individual
       absolute percentage errors weighted by the true value.

    Args:
        y_true (TensorLike): The true values of the time series.
        y_pred (TensorLike): The forecasted values of the time series.
        mask (TensorLike): Mask to be applied on the time series. Defaults to None.
        axis (int, optional): Axis along which metrics are computed. Defaults to 0.
        name (str, optional): Python str name prefixed to ops created by this function.

    Returns:
        tf.Tensor: The weighted mean absolute percentage error calculated along `axis`.
    """
    with tf.name_scope(name):
        return tf.reduce_sum(tf.abs(residuals(y_true, y_pred, mask)), axis=axis) / (
            tf.reduce_sum(mask * y_true, axis=axis) + EPSILON
        )