Mean metric

@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"),
            )
        )