# Copyright (c) 2021 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.

import paddle.fluid.ir as ir


def set_resnet_unit_attrs(resnet_unit, has_shortcut):
    resnet_unit.SetAttr("fuse_add", False)
    resnet_unit.SetAttr("act_type", "relu")
    resnet_unit.SetAttr("has_shortcut", has_shortcut)
    resnet_unit.SetAttr("data_format", 'NHWC')
    resnet_unit.SetAttr("dilation", 1)
    resnet_unit.Attr("stride").MappedPattern(op="conv2d",
                                             name="strides",
                                             element_index=0)
    resnet_unit.Attr("padding").MappedPattern(op="conv2d",
                                              name="paddings",
                                              element_index=0)
    resnet_unit.Attr("group").MappedPattern(op="conv2d", name="groups")
    resnet_unit.Attr("op_device").MappedPattern(op="conv2d", name="op_device")
    resnet_unit.Attr("op_namescope").MappedPattern(op="conv2d",
                                                   name="op_namescope")
    resnet_unit.Attr("momentum").MappedPattern(op="batch_norm", name="momentum")
    resnet_unit.Attr("epsilon").MappedPattern(op="batch_norm", name="epsilon")
    resnet_unit.Attr("use_global_stats").MappedPattern(op="batch_norm",
                                                       name="use_global_stats")


def set_resnet_unit_outputs(resnet_unit, meanX, varX, meanZ=None, varZ=None):
    resnet_unit.SetOutputs(RunningMeanX=meanX,
                           RunningVarX=varX,
                           RunningMeanZ=meanZ,
                           RunningVarZ=varZ)


@ir.RegisterPass
def fuse_resnet_unit():

    def pattern_conv_bn(x, filter, scale, bias, mean, var):
        filter.Attr("shape")[0].Mod(32).EQ(0)
        filter.Attr("shape")[1].Mod(8).EQ(0)
        filter.Attr("shape")[2].EQ(1)
        filter.Attr("shape")[3].EQ(1)
        conv2d = ir.PassDesc.OP.conv2d(Input=x, Filter=filter)
        conv2d.SetAttr("data_format", 'NHWC')
        bn = ir.PassDesc.OP.batch_norm(X=conv2d,
                                       Bias=bias,
                                       Mean=mean,
                                       Scale=scale,
                                       Variance=var)
        return bn

    def pattern_one_input(x, filter, scale, bias, mean, var):
        bn = pattern_conv_bn(x, filter, scale, bias, mean, var)
        relu = ir.PassDesc.OP.relu(X=bn.Output("Y"))
        return relu

    def replace_one_input(x, filter, scale, bias, mean, var):
        resnet_unit = ir.PassDesc.OP.resnet_unit(X=x,
                                                 FilterX=filter,
                                                 ScaleX=scale,
                                                 BiasX=bias,
                                                 MeanX=mean,
                                                 VarX=var)
        set_resnet_unit_attrs(resnet_unit, False)
        set_resnet_unit_outputs(resnet_unit, mean, var)
        return resnet_unit.Output("Y")

    def pattern_two_input(x, filterX, scaleX, biasX, meanX, varX, z, filterZ,
                          scaleZ, biasZ, meanZ, varZ):
        bnX = pattern_conv_bn(x, filterX, scaleX, biasX, meanX, varX)
        bnZ = pattern_conv_bn(x, filterZ, scaleZ, biasZ, meanZ, varZ)
        ewadd = ir.PassDesc.OP.elementwise_add(X=bnX.Output("Y"),
                                               Y=bnZ.Output("Y"))
        relu = ir.PassDesc.OP.relu(X=ewadd)
        return relu

    def replace_two_input(x, filterX, scaleX, biasX, meanX, varX, z, filterZ,
                          scaleZ, biasZ, meanZ, varZ):
        resnet_unit = ir.PassDesc.OP.resnet_unit(X=x,
                                                 FilterX=filterX,
                                                 ScaleX=scaleX,
                                                 BiasX=biasX,
                                                 MeanX=meanX,
                                                 VarX=varX,
                                                 Z=z,
                                                 FilterZ=filterZ,
                                                 ScaleZ=scaleZ,
                                                 BiasZ=biasZ,
                                                 MeanZ=meanZ,
                                                 VarZ=varZ)
        set_resnet_unit_attrs(resnet_unit, True)
        set_resnet_unit_outputs(resnet_unit, meanX, varX, meanZ, varZ)
        return resnet_unit.Output("Y")

    return (pattern_one_input, replace_one_input), (pattern_two_input,
                                                    replace_two_input)