Metric wrappers and reduction metrics

@synalinks_export("synalinks.metrics.Mean")
class Mean(Metric):
    """Compute the mean of the given values.

    For example, if values is `[1, 3, 5, 7]` then the mean is 4.

    This metric creates two variables, `total` and `count`.
    The mean value returned is simply `total` divided by `count`.


    Compilation example:

    ```python
    program.compile(
        metrics=[
            synalinks.metrics.Mean(),
        ],
    )
    ```

    Args:
        name (str): (Optional) string name of the metric instance.
        in_mask (list): (Optional) list of keys to keep to compute the metric.
        out_mask (list): (Optional) list of keys to remove to compute the metric.
        in_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are kept (combined with ``in_mask`` via OR).
        out_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are dropped (combined with ``out_mask`` via OR).

    Example:

    ```python
    >>> m = Mean()
    >>> m.update_state([1, 3, 5, 7])
    >>> m.result()
    4.0
    ```
    """

    def __init__(
        self,
        name="mean",
        in_mask=None,
        out_mask=None,
        in_mask_pattern=None,
        out_mask_pattern=None,
    ):
        super().__init__(
            name=name,
            in_mask=in_mask,
            out_mask=out_mask,
            in_mask_pattern=in_mask_pattern,
            out_mask_pattern=out_mask_pattern,
        )
        self.total_with_count = self.add_variable(
            data_model=TotalWithCount, name="total_with_count"
        )

    async def update_state(self, values):
        values = reduce_to_samplewise_values(values, reduce_fn=numpy.mean)
        total = self.total_with_count.get("total")
        self.total_with_count.update({"total": float(total + numpy.sum(values))})
        if len(values.shape) >= 1:
            num_samples = numpy.shape(values)[0]
        else:
            num_samples = 1
        count = self.total_with_count.get("count")
        self.total_with_count.update({"count": int(count + num_samples)})

    def reset_state(self):
        self.total_with_count.assign(TotalWithCount())

    def result(self):
        return float(
            numpy.divide_no_nan(
                self.total_with_count.get("total"),
                self.total_with_count.get("count"),
            )
        )

@synalinks_export("synalinks.metrics.MeanMetricWrapper")
class MeanMetricWrapper(Mean):
    """Wrap a stateless metric function with the `Mean` metric.

    You could use this class to quickly build a mean metric from a function. The
    function needs to have the signature `fn(y_true, y_pred)` and return a
    per-sample reward array. `MeanMetricWrapper.result()` will return
    the average metric value across all samples seen so far.


    Example:

    ```python
    program.compile(
        metrics=[
            synalinks.metrics.MeanMetricWrapper(),
        ],
    )
    ```

    Args:
        fn (callable): The metric function to wrap, with signature
            `fn(y_true, y_pred, **kwargs)`.
        name (str): (Optional) string name of the metric instance.
        in_mask (list): (Optional) list of keys to keep to compute the metric.
        out_mask (list): (Optional) list of keys to remove to compute the metric.
        in_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are kept (combined with ``in_mask`` via OR).
        out_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are dropped (combined with ``out_mask`` via OR).
        **kwargs (keyword arguments): Keyword arguments to pass on to `fn`.
    """

    def __init__(
        self,
        fn,
        name=None,
        in_mask=None,
        out_mask=None,
        in_mask_pattern=None,
        out_mask_pattern=None,
        **kwargs,
    ):
        super().__init__(
            name=name,
            in_mask=in_mask,
            out_mask=out_mask,
            in_mask_pattern=in_mask_pattern,
            out_mask_pattern=out_mask_pattern,
        )
        self._fn = fn
        self._fn_kwargs = kwargs

        # If we are wrapping a Synalinks reward, register the metric's
        # direction as "up" (needs to be maximized during training).
        if (
            self._fn in rewards.ALL_OBJECTS
            or hasattr(self._fn, "__class__")
            and self._fn.__class__ in rewards.ALL_OBJECTS
        ):
            self.direction = "up"

    async def update_state(self, y_true, y_pred):
        y_pred = tree.map_structure(lambda x: ops.convert_to_json_data_model(x), y_pred)
        y_true = tree.map_structure(lambda x: ops.convert_to_json_data_model(x), y_true)
        if self.in_mask or self.in_mask_pattern:
            y_pred = tree.map_structure(
                lambda x: (
                    x.in_mask(mask=self.in_mask, pattern=self.in_mask_pattern)
                    if x is not None
                    else x
                ),
                y_pred,
            )
            y_true = tree.map_structure(
                lambda x: (
                    x.in_mask(mask=self.in_mask, pattern=self.in_mask_pattern)
                    if x is not None
                    else x
                ),
                y_true,
            )
        if self.out_mask or self.out_mask_pattern:
            y_pred = tree.map_structure(
                lambda x: (
                    x.out_mask(mask=self.out_mask, pattern=self.out_mask_pattern)
                    if x is not None
                    else x
                ),
                y_pred,
            )
            y_true = tree.map_structure(
                lambda x: (
                    x.out_mask(mask=self.out_mask, pattern=self.out_mask_pattern)
                    if x is not None
                    else x
                ),
                y_true,
            )
        values = await self._fn(y_true, y_pred, **self._fn_kwargs)
        return await super().update_state(values)

    def get_config(self):
        """Returns the serializable config of the metric."""
        base_config = super().get_config()
        config = {
            "fn": serialization_lib.serialize_synalinks_object(self._fn),
        }
        config.update(self._fn_kwargs)
        return {**base_config, **config}

    @classmethod
    def from_config(cls, config):
        fn = serialization_lib.deserialize_synalinks_object(config.pop("fn"))
        return cls(fn=fn, **config)

class RewardTracker(Metric):
    """Track per-sample rewards and report them with the reward's reduction.

    Statistics are kept over every sample seen since the last reset, so the
    result does not depend on how the samples were split into batches.
    ``"none"``/``None`` report the mean, like the other scalar reward consumers.

    Args:
        reduction (str): One of ``"mean"``, ``"sum"``, ``"min"``, ``"max"``,
            ``"none"`` or ``None``.
        name (str): Optional name of the tracker.
    """

    def __init__(self, reduction="mean", name="reward"):
        super().__init__(name=name)
        self.reduction = reduction
        self.stats = self.add_variable(data_model=RewardStats, name="total_with_count")

    async def update_state(self, values):
        values = numpy.convert_to_tensor(values).ravel()
        if not values.size:
            return
        current_min = self.stats.get("min")
        current_max = self.stats.get("max")
        batch_min = float(numpy.min(values))
        batch_max = float(numpy.max(values))
        self.stats.update(
            {
                "total": float(self.stats.get("total") + numpy.sum(values)),
                "count": int(self.stats.get("count") + values.size),
                "min": batch_min if current_min is None else min(current_min, batch_min),
                "max": batch_max if current_max is None else max(current_max, batch_max),
            }
        )

    def reset_state(self):
        self.stats.assign(RewardStats())

    def result(self):
        if not self.stats.get("count"):
            return 0.0
        if self.reduction == "sum":
            return float(self.stats.get("total"))
        if self.reduction in ("min", "max"):
            return float(self.stats.get(self.reduction))
        return float(
            numpy.divide_no_nan(self.stats.get("total"), self.stats.get("count"))
        )

@synalinks_export("synalinks.metrics.Sum")
class Sum(Metric):
    """Compute the (weighted) sum of the given values.

    For example, if `values` is `[1, 3, 5, 7]` then their sum is 16.

    This metric creates one variable, `total`.
    This is ultimately returned as the sum value.


    Compilation example:

    ```python
    program.compile(
        metrics=[
            synalinks.metrics.Sum(),
        ],
    )
    ```

    Args:
        name (str): (Optional) string name of the metric instance.
        in_mask (list): (Optional) list of keys to keep to compute the metric.
        out_mask (list): (Optional) list of keys to remove to compute the metric.
        in_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are kept (combined with ``in_mask`` via OR).
        out_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are dropped (combined with ``out_mask`` via OR).

    Example:

    ```python
    >>> m = metrics.Sum()
    >>> m.update_state([1, 3, 5, 7])
    >>> m.result()
    16.0
    ```
    """

    def __init__(
        self,
        name="sum",
        in_mask=None,
        out_mask=None,
        in_mask_pattern=None,
        out_mask_pattern=None,
    ):
        super().__init__(
            name=name,
            in_mask=in_mask,
            out_mask=out_mask,
            in_mask_pattern=in_mask_pattern,
            out_mask_pattern=out_mask_pattern,
        )
        self.total = self.add_variable(
            data_model=Total,
            name="total",
        )

    async def update_state(self, values):
        values = reduce_to_samplewise_values(values, reduce_fn=numpy.sum)
        total = self.total.get("total")
        self.total.update({"total": float(numpy.sum(total, values))})

    def reset_state(self):
        self.total.assign(Total())

    def result(self):
        return self.total.get("total")