from .layer import Layer
from .convolution import *
from .connected import *


class avgpool_layer(Layer):
    pass


class crop_layer(Layer):
    pass


class maxpool_layer(Layer):
    def setup(self, ksize, stride, pad):
        self.stride = stride
        self.ksize = ksize
        self.pad = pad


class softmax_layer(Layer):
    def setup(self, groups):
        self.groups = groups


class dropout_layer(Layer):
    def setup(self, p):
        self.h['pdrop'] = dict({
            'feed': p,  # for training
            'dfault': 1.0,  # for testing
            'shape': ()
        })


class route_layer(Layer):
    def setup(self, routes):
        self.routes = routes


class reorg_layer(Layer):
    def setup(self, stride):
        self.stride = stride


"""
Darkop Factory
"""

darkops = {
    'dropout': dropout_layer,
    'connected': connected_layer,
    'maxpool': maxpool_layer,
    'convolutional': convolutional_layer,
    'avgpool': avgpool_layer,
    'softmax': softmax_layer,
    'crop': crop_layer,
    'local': local_layer,
    'select': select_layer,
    'route': route_layer,
    'reorg': reorg_layer,
    'conv-select': conv_select_layer,
    'conv-extract': conv_extract_layer,
    'extract': extract_layer
}


def create_darkop(ltype, num, *args):
    op_class = darkops.get(ltype, Layer)
    return op_class(ltype, num, *args)