r"""
    Batch normalization.

    .. math::
        \begin{eqnarray}
          \mu &=& \frac{1}{M} \sum x_i \\
          \sigma^2 &=& \frac{1}{M} \sum \left(x_i - \mu\right)^2 \\
          \hat{x}_i &=& \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \\
          y_i &=& \hat{x}_i \gamma + \beta.
        \end{eqnarray}


    At testing time, the mean and variance values used are those that were computed during training by moving average.

    References:

        * `Ioffe and Szegedy, Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.
          <https://arxiv.org/abs/1502.03167>`_

    Args:
        x(~nnabla.Variable): N-D array of input.
        beta(~nnabla.Variable): N-D array of beta which is learned.
        gamma(~nnabla.Variable): N-D array of gamma which is learned.
        mean(~nnabla.Variable): N-D array of running mean (modified during forward execution).
        variance(~nnabla.Variable): N-D array of running variance (modified during forward execution).
        axes(repeated int64): Axes mean and variance are taken.
        decay_rate(float): Decay rate of running mean and variance.
        eps(float): Tiny value to avoid zero division by std.
        batch_stat(bool): Use mini-batch statistics rather than running ones.
        output_stat(bool): It true, the batch statistics of mean and variance,
            will be returned as Variables. They are also differentiable.

    Returns:
        Returns batch normalization output as :obj:`~nnabla.Variable`.
        If ``output_stat=True``, it also returns the mean and variance
        of the mini-batch

        * :obj:`~nnabla.Variable`: Output of the batch normalization
        * :obj:`~nnabla.Variable`: Mean (if ``output_stat=True`)
        * :obj:`~nnabla.Variable`: Variance (if ``output_stat=True`)

    See Also:
        ``nnabla.function_bases.batch_normalization``.

def batch_normalization(x, beta, gamma, mean, variance, axes=[1], decay_rate=0.9, eps=1e-05, batch_stat=True, output_stat=False, n_outputs=None):
    r"""
    Batch normalization.

    .. math::
        \begin{eqnarray}
          \mu &=& \frac{1}{M} \sum x_i \\
          \sigma^2 &=& \frac{1}{M} \sum \left(x_i - \mu\right)^2 \\
          \hat{x}_i &=& \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \\
          y_i &=& \hat{x}_i \gamma + \beta.
        \end{eqnarray}


    At testing time, the mean and variance values used are those that were computed during training by moving average.

    References:

        * `Ioffe and Szegedy, Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.
          <https://arxiv.org/abs/1502.03167>`_

    Args:
        x(~nnabla.Variable): N-D array of input.
        beta(~nnabla.Variable): N-D array of beta which is learned.
        gamma(~nnabla.Variable): N-D array of gamma which is learned.
        mean(~nnabla.Variable): N-D array of running mean (modified during forward execution).
        variance(~nnabla.Variable): N-D array of running variance (modified during forward execution).
        axes(repeated int64): Axes mean and variance are taken.
        decay_rate(float): Decay rate of running mean and variance.
        eps(float): Tiny value to avoid zero division by std.
        batch_stat(bool): Use mini-batch statistics rather than running ones.
        output_stat(bool): It true, the batch statistics of mean and variance,
            will be returned as Variables. They are also differentiable.

    Returns:
        Returns batch normalization output as :obj:`~nnabla.Variable`.
        If ``output_stat=True``, it also returns the mean and variance
        of the mini-batch

        * :obj:`~nnabla.Variable`: Output of the batch normalization
        * :obj:`~nnabla.Variable`: Mean (if ``output_stat=True`)
        * :obj:`~nnabla.Variable`: Variance (if ``output_stat=True`)

    See Also:
        ``nnabla.function_bases.batch_normalization``.

    """
    from .function_bases import batch_normalization as batch_normalization_base
    n_outputs = 3 if output_stat else 1
    assert batch_stat or (not output_stat)
    if batch_stat and (mean.parent or variance.parent) is not None:
        raise ValueError(
            "if batch_stat is True, mean and variable must not have a parent function")

    if len(axes) == 1:
        return batch_normalization_base(x, beta, gamma, mean, variance,
                                        axes=axes,
                                        decay_rate=decay_rate,
                                        eps=eps,
                                        batch_stat=batch_stat,
                                        n_outputs=n_outputs)

    def transpose_and_reshape(x, axes):
        transposed = transpose(x, transpose_axes)
        return reshape(transposed, [rd(lambda x, y: x * y, transposed.shape[:len(axes)])] + list(
            transposed.shape[len(axes):])), transposed.shape

    def inverse_transpose_and_reshape(x, axes, variable_shape):
        un_reshaped = reshape(
            x, list(variable_shape[:len(axes)] + variable_shape[len(axes):]))
        return transpose(un_reshaped, inv_transpose_axes)

    def get_tranpose_args(ndim, axes):
        transpose_axes = [i for i in list(
            axes)] + [i for i in range(ndim) if i not in list(axes)]
        inv_transpose_axes = np.argsort(transpose_axes).tolist()
        return transpose_axes, inv_transpose_axes

    transpose_axes, inv_transpose_axes = get_tranpose_args(len(x.shape), axes)
    inp, transposed_inp_shape = transpose_and_reshape(x, axes)
    beta, transposed_beta_shape = transpose_and_reshape(beta, axes)
    gamma, transposed_gamma_shape = transpose_and_reshape(gamma, axes)
    mean, transposed_mean_shape = transpose_and_reshape(mean, axes)
    variance, transposed_variance_shape = transpose_and_reshape(variance, axes)

    if n_outputs == 1:
        out = batch_normalization_base(inp, beta, gamma, mean, variance,
                                       axes=[0],
                                       decay_rate=decay_rate,
                                       eps=eps,
                                       batch_stat=batch_stat,
                                       n_outputs=n_outputs)
        return inverse_transpose_and_reshape(out, axes, transposed_inp_shape)
    out, mean, variance = batch_normalization_base(inp, beta, gamma, mean, variance,
                                                   axes=[0],
                                                   decay_rate=decay_rate,
                                                   eps=eps,
                                                   batch_stat=batch_stat,
                                                   n_outputs=n_outputs)
    out = inverse_transpose_and_reshape(out, axes, transposed_inp_shape)
    mean = inverse_transpose_and_reshape(mean, axes, transposed_mean_shape)
    variance = inverse_transpose_and_reshape(
        variance, axes, transposed_variance_shape)
    return out, mean, variance