from __future__ import print_function
### generated by caffe2fluid, your net is in class "LeNet" ###

from builtins import str
from past.builtins import basestring
from builtins import object
import sys
import os
import math
import numpy as np


def import_fluid():
    import paddle.fluid as fluid
    return fluid


def layer(op):
    '''Decorator for composable network layers.'''

    def layer_decorated(self, *args, **kwargs):
        # Automatically set a name if not provided.
        name = kwargs.setdefault('name', self.get_unique_name(op.__name__))
        # Figure out the layer inputs.
        if len(self.terminals) == 0:
            raise RuntimeError('No input variables found for layer %s.' % name)
        elif len(self.terminals) == 1:
            layer_input = self.terminals[0]
        else:
            layer_input = list(self.terminals)

        # Perform the operation and get the output.
        layer_output = op(self, layer_input, *args, **kwargs)
        # Add to layer LUT.
        self.layers[name] = layer_output
        # This output is now the input for the next layer.
        self.feed(layer_output)
        # Return self for chained calls.
        return self

    return layer_decorated


class Network(object):
    def __init__(self, inputs, trainable=True):
        # The input nodes for this network
        self.inputs = inputs
        # The current list of terminal nodes
        self.terminals = []
        # Mapping from layer names to layers
        self.layers = dict(inputs)
        # If true, the resulting variables are set as trainable
        self.trainable = trainable
        # Switch variable for dropout
        self.paddle_env = None
        self.output_names = []
        self.name_trace = None
        self.setup()

    def setup(self):
        '''Construct the network. '''
        raise NotImplementedError('Must be implemented by the subclass.')

    def load(self, data_path, exe=None, place=None, ignore_missing=False):
        '''Load network weights.
        data_path: The path to the numpy-serialized network weights
        ignore_missing: If true, serialized weights for missing layers are ignored.
        '''
        fluid = import_fluid()
        #load fluid mode directly
        if os.path.isdir(data_path):
            assert (exe is not None), \
                'must provide a executor to load fluid model'
            fluid.io.load_persistables(executor=exe, dirname=data_path)
            return True

        #load model from a npy file
        if exe is None or place is None:
            if self.paddle_env is None:
                place = fluid.CPUPlace()
                exe = fluid.Executor(place)
                self.paddle_env = {'place': place, 'exe': exe}
                exe = exe.run(fluid.default_startup_program())
            else:
                place = self.paddle_env['place']
                exe = self.paddle_env['exe']

        data_dict = np.load(data_path).item()
        for op_name in data_dict:
            if op_name == 'caffe2fluid_name_trace':
                self.name_trace = data_dict[op_name]
                continue

            layer = self.layers[op_name]
            for param_name, data in list(data_dict[op_name].items()):
                try:
                    name = '%s_%s' % (op_name, param_name)
                    v = fluid.global_scope().find_var(name)
                    w = v.get_tensor()
                    w.set(data.reshape(w.shape()), place)
                except ValueError:
                    if not ignore_missing:
                        raise
        return True

    def feed(self, *args):
        '''Set the input(s) for the next operation by replacing the terminal nodes.
        The arguments can be either layer names or the actual layers.
        '''
        assert len(args) != 0
        self.terminals = []
        for fed_layer in args:
            if isinstance(fed_layer, basestring):
                try:
                    fed_layer = self.layers[fed_layer]
                except KeyError:
                    raise KeyError('Unknown layer name fed: %s' % fed_layer)
            self.terminals.append(fed_layer)
        return self

    def get_output(self):
        '''Returns the current network output.'''
        return self.terminals[-1]

    def get_unique_name(self, prefix):
        '''Returns an index-suffixed unique name for the given prefix.
        This is used for auto-generating layer names based on the type-prefix.
        '''
        ident = sum(t.startswith(prefix) for t, _ in list(self.layers.items())) + 1
        return '%s_%d' % (prefix, ident)

    def get_unique_output_name(self, prefix, layertype):
        '''Returns an index-suffixed unique name for the given prefix.
            This is used for auto-generating layer names based on the type-prefix.
        '''
        ident = sum(t.startswith(prefix) for t in self.output_names) + 1
        unique_name = '%s.%s.output.%d' % (prefix, layertype, ident)
        self.output_names.append(unique_name)
        return unique_name

    @layer
    def conv(self,
             input,
             k_h,
             k_w,
             c_o,
             s_h,
             s_w,
             name,
             relu=True,
             relu_negative_slope=0.0,
             padding=None,
             dilation=1,
             group=1,
             biased=True):
        if padding is None:
            padding = [0, 0]

        # Get the number of channels in the input
        c_i, h_i, w_i = input.shape[1:]

        # Verify that the grouping parameter is valid
        assert c_i % group == 0
        assert c_o % group == 0

        fluid = import_fluid()
        prefix = name + '_'
        leaky_relu = False
        act = 'relu'
        if relu is False:
            act = None
        elif relu_negative_slope != 0.0:
            leaky_relu = True
            act = None

        output = fluid.layers.conv2d(
            name=self.get_unique_output_name(name, 'conv2d'),
            input=input,
            filter_size=[k_h, k_w],
            num_filters=c_o,
            stride=[s_h, s_w],
            padding=padding,
            dilation=dilation,
            groups=group,
            param_attr=fluid.ParamAttr(name=prefix + "weights"),
            bias_attr=fluid.ParamAttr(name=prefix + "biases"),
            act=act)

        if leaky_relu:
            output = fluid.layers.leaky_relu(output, alpha=relu_negative_slope)

        return output

    @layer
    def deconv(self,
               input,
               k_h,
               k_w,
               c_o,
               s_h,
               s_w,
               name,
               relu=True,
               relu_negative_slope=0.0,
               padding=None,
               dilation=1,
               biased=True):
        if padding is None:
            padding = [0, 0]

        # Get the number of channels in the input
        c_i, h_i, w_i = input.shape[1:]

        fluid = import_fluid()
        prefix = name + '_'
        leaky_relu = False
        act = 'relu'
        if relu is False:
            act = None
        elif relu_negative_slope != 0.0:
            leaky_relu = True
            act = None

        p_h = padding[0]
        p_w = padding[1]
        h_o = (h_i - 1) * s_h - 2 * p_h + dilation * (k_h - 1) + 1
        w_o = (w_i - 1) * s_w - 2 * p_w + dilation * (k_w - 1) + 1
        output = fluid.layers.conv2d_transpose(
            name=self.get_unique_output_name(name, 'conv2d_transpose'),
            input=input,
            num_filters=c_o,
            output_size=[h_o, w_o],
            filter_size=[k_h, k_w],
            padding=padding,
            stride=[s_h, s_w],
            dilation=dilation,
            param_attr=fluid.ParamAttr(name=prefix + "weights"),
            bias_attr=fluid.ParamAttr(name=prefix + "biases"),
            act=act)

        if leaky_relu:
            output = fluid.layers.leaky_relu(output, alpha=relu_negative_slope)

        return output

    @layer
    def relu(self, input, name):
        fluid = import_fluid()
        output = fluid.layers.relu(input)
        return output

    @layer
    def prelu(self, input, channel_shared, name):
        fluid = import_fluid()
        if channel_shared:
            mode = 'all'
        else:
            mode = 'channel'

        prefix = name + '_'
        output = fluid.layers.prelu(
            input,
            mode=mode,
            param_attr=fluid.ParamAttr(name=prefix + 'negslope'))
        return output

    def pool(self, pool_type, input, k_h, k_w, s_h, s_w, ceil_mode, padding,
             name):
        # Get the number of channels in the input
        in_hw = input.shape[2:]
        k_hw = [k_h, k_w]
        s_hw = [s_h, s_w]

        fluid = import_fluid()
        output = fluid.layers.pool2d(
            name=name,
            input=input,
            pool_size=k_hw,
            pool_stride=s_hw,
            pool_padding=padding,
            ceil_mode=ceil_mode,
            pool_type=pool_type)
        return output

    @layer
    def max_pool(self,
                 input,
                 k_h,
                 k_w,
                 s_h,
                 s_w,
                 ceil_mode,
                 padding=[0, 0],
                 name=None):
        return self.pool(
            'max',
            input,
            k_h,
            k_w,
            s_h,
            s_w,
            ceil_mode,
            padding,
            name=self.get_unique_output_name(name, 'max_pool'))

    @layer
    def avg_pool(self,
                 input,
                 k_h,
                 k_w,
                 s_h,
                 s_w,
                 ceil_mode,
                 padding=[0, 0],
                 name=None):
        return self.pool(
            'avg',
            input,
            k_h,
            k_w,
            s_h,
            s_w,
            ceil_mode,
            padding,
            name=self.get_unique_output_name(name, 'avg_pool'))

    @layer
    def sigmoid(self, input, name):
        fluid = import_fluid()
        return fluid.layers.sigmoid(
            input, name=self.get_unique_output_name(name, 'sigmoid'))

    @layer
    def tanh(self, input, name):
        fluid = import_fluid()
        return fluid.layers.tanh(
            input, name=self.get_unique_output_name(name, 'tanh'))

    @layer
    def lrn(self, input, radius, alpha, beta, name, bias=1.0):
        fluid = import_fluid()
        output = fluid.layers.lrn(input=input,
                                  n=radius,
                                  k=bias,
                                  alpha=alpha,
                                  beta=beta,
                                  name=self.get_unique_output_name(name, 'lrn'))
        return output

    @layer
    def concat(self, inputs, axis, name):
        fluid = import_fluid()
        output = fluid.layers.concat(
            input=inputs,
            axis=axis,
            name=self.get_unique_output_name(name, 'concat'))
        return output

    @layer
    def add(self, inputs, name):
        fluid = import_fluid()
        output = inputs[0]
        for i in inputs[1:]:
            output = fluid.layers.elementwise_add(
                x=output, y=i, name=self.get_unique_output_name(name, 'add'))
        return output

    @layer
    def max(self, inputs, name):
        fluid = import_fluid()
        output = inputs[0]
        for i in inputs[1:]:
            output = fluid.layers.elementwise_max(
                x=output, y=i, name=self.get_unique_output_name(name, 'max'))
        return output

    @layer
    def multiply(self, inputs, name):
        fluid = import_fluid()
        output = inputs[0]
        for i in inputs[1:]:
            output = fluid.layers.elementwise_mul(
                x=output, y=i, name=self.get_unique_output_name(name, 'mul'))
        return output

    @layer
    def fc(self, input, num_out, name, relu=True, act=None):
        fluid = import_fluid()

        if act is None:
            act = 'relu' if relu is True else None

        prefix = name + '_'
        output = fluid.layers.fc(
            name=self.get_unique_output_name(name, 'fc'),
            input=input,
            size=num_out,
            act=act,
            param_attr=fluid.ParamAttr(name=prefix + 'weights'),
            bias_attr=fluid.ParamAttr(name=prefix + 'biases'))
        return output

    @layer
    def softmax(self, input, name):
        fluid = import_fluid()
        shape = input.shape
        if len(shape) > 2:
            for sz in shape[2:]:
                assert sz == 1, "invalid input shape[%s] for softmax" % (
                    str(shape))
            input = fluid.layers.reshape(input, shape[0:2])

        output = fluid.layers.softmax(
            input, name=self.get_unique_output_name(name, 'softmax'))
        return output

    @layer
    def batch_normalization(self,
                            input,
                            name,
                            scale_offset=True,
                            eps=1e-5,
                            relu=False,
                            relu_negative_slope=0.0):
        # NOTE: Currently, only inference is supported
        fluid = import_fluid()
        prefix = name + '_'
        param_attr = None if scale_offset is False else fluid.ParamAttr(
            name=prefix + 'scale')
        bias_attr = None if scale_offset is False else fluid.ParamAttr(
            name=prefix + 'offset')
        mean_name = prefix + 'mean'
        variance_name = prefix + 'variance'

        leaky_relu = False
        act = 'relu'
        if relu is False:
            act = None
        elif relu_negative_slope != 0.0:
            leaky_relu = True
            act = None

        output = fluid.layers.batch_norm(
            name=self.get_unique_output_name(name, 'batch_norm'),
            input=input,
            is_test=True,
            param_attr=param_attr,
            bias_attr=bias_attr,
            moving_mean_name=mean_name,
            moving_variance_name=variance_name,
            epsilon=eps,
            act=act)

        if leaky_relu:
            output = fluid.layers.leaky_relu(output, alpha=relu_negative_slope)

        return output

    @layer
    def dropout(self, input, drop_prob, name, is_test=True):
        fluid = import_fluid()
        if is_test:
            output = input
        else:
            output = fluid.layers.dropout(
                input,
                dropout_prob=drop_prob,
                is_test=is_test,
                name=self.get_unique_output_name(name, 'dropout'))
        return output

    @layer
    def scale(self, input, axis=1, num_axes=1, name=None):
        fluid = import_fluid()

        assert num_axes == 1, "layer scale not support this num_axes[%d] now" % (
            num_axes)

        prefix = name + '_'
        scale_shape = input.shape[axis:axis + num_axes]
        param_attr = fluid.ParamAttr(name=prefix + 'scale')
        scale_param = fluid.layers.create_parameter(
            shape=scale_shape, dtype=input.dtype, name=name, attr=param_attr)

        offset_attr = fluid.ParamAttr(name=prefix + 'offset')
        offset_param = fluid.layers.create_parameter(
            shape=scale_shape, dtype=input.dtype, name=name, attr=offset_attr)

        output = fluid.layers.elementwise_mul(
            input,
            scale_param,
            axis=axis,
            name=self.get_unique_output_name(name, 'scale_mul'))
        output = fluid.layers.elementwise_add(
            output,
            offset_param,
            axis=axis,
            name=self.get_unique_output_name(name, 'scale_add'))
        return output

    def custom_layer_factory(self):
        """ get a custom layer maker provided by subclass
        """
        raise NotImplementedError(
            '[custom_layer_factory] must be implemented by the subclass.')

    @layer
    def custom_layer(self, inputs, kind, name, *args, **kwargs):
        """ make custom layer
        """
        name = self.get_unique_output_name(name, kind)
        layer_factory = self.custom_layer_factory()
        return layer_factory(kind, inputs, name, *args, **kwargs)


class LeNet(Network):
    ### automatically generated by caffe2fluid ###
    inputs_info = {"data": [1, 28, 28]}
    custom_layers_path = "/mnt/advbox-demo/caffe2fluid/kaffe"

    def custom_layer_factory(self):
        import os

        pk_paths = []
        default = os.path.dirname(os.path.abspath(__file__))
        location = os.environ.get('CAFFE2FLUID_CUSTOM_LAYERS', default)
        pk_name = 'custom_layers'
        pk_dir = os.path.join(location, pk_name)
        pk_paths.append((location, pk_dir))

        location = LeNet.custom_layers_path
        pk_dir = os.path.join(LeNet.custom_layers_path, pk_name)
        pk_paths.append((location, pk_dir))

        for loc, pk_dir in pk_paths:
            if os.path.exists(pk_dir):
                if loc not in sys.path:
                    sys.path.insert(0, loc)
                    break

        try:
            from custom_layers import make_custom_layer
            return make_custom_layer
        except Exception as e:
            print('maybe you should set $CAFFE2FLUID_CUSTOM_LAYERS first')
            raise e

    @classmethod
    def input_shapes(cls):
        return cls.inputs_info

    @classmethod
    def convert(cls, npy_model, fluid_path, outputs=None):
        fluid = import_fluid()
        shapes = cls.input_shapes()
        input_name = list(shapes.keys())[0]
        feed_data = {}
        for name, shape in list(shapes.items()):
            data_layer = fluid.layers.data(
                name=name, shape=shape, dtype="float32")
            feed_data[name] = data_layer

        net = cls(feed_data)
        place = fluid.CPUPlace()
        exe = fluid.Executor(place)
        exe.run(fluid.default_startup_program())
        net.load(data_path=npy_model, exe=exe, place=place)
        output_vars = []

        model_filename = 'model'
        params_filename = 'params'
        if outputs is None:
            output_vars.append(net.get_output())
        else:
            if outputs[0] == 'dump_all':
                model_filename = None
                params_filename = None
                output_vars.append(net.get_output())
            else:
                if type(outputs) is list:
                    for n in outputs:
                        assert n in net.layers, 'not found layer with this name[%s]' % (
                            n)
                        output_vars.append(net.layers[n])

        fluid.io.save_inference_model(
            fluid_path, [input_name],
            output_vars,
            exe,
            main_program=None,
            model_filename=model_filename,
            params_filename=params_filename)
        return 0

    def setup(self):
        self.feed('data')
        self.conv(5, 5, 20, 1, 1, relu=False, name='conv1')
        self.max_pool(2, 2, 2, 2, True, name='pool1')
        self.conv(5, 5, 50, 1, 1, relu=False, name='conv2')
        self.max_pool(2, 2, 2, 2, True, name='pool2')
        self.fc(500, name='ip1')
        self.fc(10, relu=False, name='ip2')
        self.softmax(name='prob')


def main():
    """ a tool used to convert caffe model to fluid
    """

    import sys
    import os
    filename = os.path.splitext(os.path.basename(sys.argv[0]))[0]
    if len(sys.argv) < 3:
        print('usage:')
        print('	python %s %s.npy [save_dir] [layer names seperated by comma]' \
                % (sys.argv[0], filename))
        print('	eg: python %s %s.npy ./fluid' % (sys.argv[0], filename))
        print('	eg: python %s %s.npy ./fluid layer_name1,layer_name2' \
                % (sys.argv[0], filename))
        return 1

    npy_weight = sys.argv[1]
    fluid_model = sys.argv[2]
    outputs = None
    if len(sys.argv) >= 4:
        outputs = sys.argv[3].split(',')

    ret = LeNet.convert(npy_weight, fluid_model, outputs)
    if ret == 0:
        outputs = 'last output layer' if outputs is None else outputs
        print('succeed to convert to fluid format with output layers[%s]'
              ' in directory[%s]' % (outputs, fluid_model))
    else:
        print('failed to convert model to fluid format')

    return ret


if __name__ == "__main__":
    exit(main())
