#  Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
#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 as fluid
from paddle.fluid.param_attr import ParamAttr
from paddle.fluid.initializer import Constant
from paddle.fluid.regularizer import L2Decay

def conv_bn_layer(input,
                  ch_out,
                  filter_size,
                  stride,
                  padding,
                  act='leaky',
                  is_test=True,
                  name=None):
    conv1 = fluid.layers.conv2d(
        input=input,
        num_filters=ch_out,
        filter_size=filter_size,
        stride=stride,
        padding=padding,
        act=None,
        param_attr=ParamAttr(initializer=fluid.initializer.Normal(0., 0.02),
                name=name+".conv.weights"),
        bias_attr=False)

    bn_name = name + ".bn"
    out = fluid.layers.batch_norm(
        input=conv1,
        act=None,
        is_test=is_test,
        param_attr=ParamAttr(
                initializer=fluid.initializer.Normal(0., 0.02),
                regularizer=L2Decay(0.),
                name=bn_name + '.scale'),
        bias_attr=ParamAttr(
                initializer=fluid.initializer.Constant(0.0),
                regularizer=L2Decay(0.),
                name=bn_name + '.offset'),
        moving_mean_name=bn_name + '.mean',
        moving_variance_name=bn_name + '.var')
    if act == 'leaky':
        out = fluid.layers.leaky_relu(x=out, alpha=0.1)
    return out

def downsample(input, ch_out, filter_size=3, stride=2, padding=1, is_test=True, name=None):
    return conv_bn_layer(input, 
            ch_out=ch_out, 
            filter_size=filter_size, 
            stride=stride, 
            padding=padding, 
            is_test=is_test,
            name=name)

def basicblock(input, ch_out, is_test=True, name=None):
    conv1 = conv_bn_layer(input, ch_out, 1, 1, 0, is_test=is_test, name=name+".0")
    conv2 = conv_bn_layer(conv1, ch_out*2, 3, 1, 1, is_test=is_test, name=name+".1")
    out = fluid.layers.elementwise_add(x=input, y=conv2, act=None)
    return out

def layer_warp(block_func, input, ch_out, count, is_test=True, name=None):
    res_out = block_func(input, ch_out, is_test=is_test, name='{}.0'.format(name))
    for j in range(1, count):
        res_out = block_func(res_out, ch_out, is_test=is_test, name='{}.{}'.format(name, j))
    return res_out

DarkNet_cfg = {
        53: ([1,2,8,8,4],basicblock)
}

def add_DarkNet53_conv_body(body_input, is_test=True):
    stages, block_func = DarkNet_cfg[53]
    stages = stages[0:5]
    conv1 = conv_bn_layer(
            body_input, ch_out=32, filter_size=3, stride=1, padding=1, is_test=is_test, name="yolo_input")
    downsample_ = downsample(conv1, ch_out=conv1.shape[1]*2, is_test=is_test, name="yolo_input.downsample")
    blocks = []
    for i, stage in enumerate(stages):
        block = layer_warp(block_func, downsample_, 32 *(2**i), stage, is_test=is_test, name="stage.{}".format(i))
        blocks.append(block)
        if i < len(stages) - 1: # do not downsaple in the last stage
            downsample_ = downsample(block, ch_out=block.shape[1]*2, is_test=is_test, name="stage.{}.downsample".format(i))
    return blocks[-1:-4:-1]

