import torch
from typing import Optional


def norm_zero_mean_unit_std(
    data: torch.Tensor, per_channel: bool = True, out: Optional[torch.Tensor] = None, eps: Optional[float] = 1e-8
) -> torch.Tensor:
    """
    Normalize mean to zero and std to one

    Args:
        data: input data. Per channel option supports [C,H,W] and [C,H,W,D].
        per_channel: range is normalized per channel
        out: if provided, result is saved in here
        eps: small constant for numerical stability.
            If None, no factor constant will be added

    Returns:
        torch.Tensor: normalized data
    """

    def _norm(_data: torch.Tensor, _out: torch.Tensor):
        denom = _data.std()
        if eps is not None:
            denom = denom + eps
        _out = (_data - _data.mean()) / denom
        return _out

    if out is None:
        out = torch.zeros_like(data)

    if per_channel:
        for _c in range(data.shape[0]):
            out[_c] = _norm(data[_c], out[_c])
    else:
        out = _norm(data, out)
    return out