# ----------------------------------------------------------------------------
# Copyright 2014 Nervana Systems Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""
Batch normalization transform functions and classes.
"""

import logging
from neon.transforms.activation import Activation
from neon.util.param import req_param, opt_param
import numpy as np


logger = logging.getLogger(__name__)


class BatchNorm(Activation):

    """
    Embodiment of a Batch Normalization Transform.  Can be used in a
    WeightLayer after the linear transform but before the typical activation
    (ReLu, sigmoid, tanh, etc.)

    The transform does z-norm along the batch dimension then adjusts using
    an affine transform with learned parameters gamma and beta.  gamma and beta
    gradients (updates) are computed via bprop, and the updates are applied
    at the same time that the associated WeightLayer is updated.  Likewise,
    the same learning rule is used for gamma and beta as the WeightLayer

    Forward pass: (gamma/beta are scalar parameters for each unit)

        x' = (x - mean) / sqrt(var + eps)
        y  = gamma * x' + beta

    Backward pass:

        dy/dx = dy/dx' * dx'/dx
        = gamma * [1*(var+eps)^-1/2 + (x-mean) * (var+eps)^-3/2 * (2x)^-1/2]

    """

    inference_warning_displayed = False

    def initialize(self, kwargs):
        """
        Initialize the Batch Normalization transform. This function will be