# 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 __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import paddle.fluid as fluid
from paddle.fluid.param_attr import ParamAttr
from .search_space_base import SearchSpaceBase
from .base_layer import conv_bn_layer
from .search_space_registry import SEARCHSPACE

__all__ = ["ResNetSpace"]


@SEARCHSPACE.register
class ResNetSpace(SearchSpaceBase):
    def __init__(self, input_size, output_size, block_num, extract_feature=False, class_dim=1000):
        super(ResNetSpace, self).__init__(input_size, output_size, block_num)
        self.filter_num1 = np.array([48,64,96,128,160,192,224]) #7 
        self.filter_num2 = np.array([64,96,128,160,192,256,320]) #7
        self.filter_num3 = np.array([128,160,192,256,320,384]) #6
        self.filter_num4 = np.array([192,256,384,512,640]) #5
        self.repeat1 = [2,3,4,5,6]#5
        self.repeat2 = [2,3,4,5,6,7]#6
        self.repeat3 = [2,3,4,5,6,7,8,10,12,14,16,20,24]#13
        self.repeat4 = [2,3,4,5,6,7]#6
        self.class_dim = class_dim
        self.extract_feature = extract_feature


    def init_tokens(self):
        init_token_base = [0,0,0,0,0,0,0,0]
        self.token_len = self.block_num * 2
        return init_token_base[:self.token_len]

    def range_table(self):
        range_table_base = [3,3,3,3,3,3,3,3]
        return range_table_base[:self.token_len]

    def token2arch(self,tokens=None):
        assert self.block_num < 5, 'block number must less than 5, but receive block number is {}'.format(self.block_num)

        if tokens is None:
            tokens = self.init_tokens()

        def net_arch(input):
            depth = []
            num_filters = []
            if self.block_num <= 1:
                filter1 = self.filter_num1[tokens[0]]
                repeat1 = self.repeat1[tokens[1]]
                depth.append(filter1)
                num_filters.append(repeat1)
            if self.block_num <= 2:
                filter2 = self.filter_num2[tokens[2]]
                repeat2 = self.repeat2[tokens[3]]
                depth.append(filter2)
                num_filters.append(repeat2)
            if self.block_num <= 3:
                filter3 = self.filter_num3[tokens[4]]
                repeat3 = self.repeat3[tokens[5]]
                depth.append(filter3)
                num_filters.append(repeat3)
            if self.block_num <= 4:
                filter4 = self.filter_num4[tokens[6]]
                repeat4 = self.repeat4[tokens[7]]
                depth.append(filter4)
                num_filters.append(repeat4)

            conv = conv_bn_layer(input=input, filter_size=5, num_filters=filter1, stride=2, act='relu')
            for block in range(len(depth)):
                for i in range(depth[block]):
                    conv = self._basicneck_block(input=conv, num_filters=num_filters[block], stride=2 if i == 0 and block != 0 else 1)

            if self.output_size == 1:
                conv = fluid.layers.fc(
                         input=conv,
                         size=self.class_dim,
                         act=None,
                         param_attr=fluid.param_attr.ParamAttr(
                             initializer=fluid.initializer.NormalInitializer(0.0,0.01)),
                         bias_attr=fluid.param_attr.ParamAttr(
                             initializer=fluid.initializer.ConstantInitializer()))

            return conv

        return net_arch

    def _shortcut(self, input, ch_out, stride):
        ch_in = input.shape[1]
        if ch_in != ch_out or stride != 1:
            return conv_bn_layer(input=input, filter_size=1, num_filters=ch_out, stride=stride)
        else:
            return input

    def _basicneck_block(self, input, num_filters, stride):
        conv0 = conv_bn_layer(input=input, filter_size=3, num_filters=num_filters, stride=stride, act='relu')
        conv1 = conv_bn_layer(input=conv0, filter_size=3, num_filters=num_filters, stride=1, act=None)
        short = self._shortcut(input, num_filters, stride)
        return fluid.layers.elementwise_add(x=short, y=conv1, act='relu')