#   Copyright (c) 2019  PaddlePaddle Authors. All Rights Reserved.
#
# 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.

from .register import register


def InstanceNormalization_shape(input_shape):
    return input_shape


def InstanceNormalization_layer(inputs, name=None):
    # TODO(lvmengsi@baidu.com): Check the accuracy when using fluid.layers.layer_norm.
    epsilon = 1e-5
    input_ = inputs[0]
    mean = fluid.layers.reduce_mean(input_, dim=[2, 3], keep_dim=True)
    var = fluid.layers.reduce_mean(fluid.layers.square(input_ - mean),
                                   dim=[2, 3],
                                   keep_dim=True)
    if name is not None:
        scale_name = name + "_scale"
        offset_name = name + "_offset"
    scale_param = fluid.ParamAttr(name=scale_name,
                                  initializer=fluid.initializer.Constant(1.0),
                                  trainable=True)
    offset_param = fluid.ParamAttr(name=offset_name,
                                   initializer=fluid.initializer.Constant(0.0),
                                   trainable=True)
    scale = fluid.layers.create_parameter(attr=scale_param,
                                          shape=input_.shape[1:2],
                                          dtype="float32")
    offset = fluid.layers.create_parameter(attr=offset_param,
                                           shape=input_.shape[1:2],
                                           dtype="float32")

    tmp = fluid.layers.elementwise_mul(x=(input_ - mean), y=scale, axis=1)
    tmp = tmp / fluid.layers.sqrt(var + epsilon)
    tmp = fluid.layers.elementwise_add(tmp, offset, axis=1)
    return tmp


def InstanceNormalization_weights(name, data=None):
    weights_name = [name + '_scale']
    return weights_name


register(kind='InstanceNormalization',
         shape=InstanceNormalization_shape,
         layer=InstanceNormalization_layer,
         child_func=None,
         weights=InstanceNormalization_weights)