/*
    Copyright (c) 2013, Taiga Nomi
    All rights reserved.

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
    notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions and the following disclaimer in the
    documentation and/or other materials provided with the distribution.
    * Neither the name of the <organization> nor the
    names of its contributors may be used to endorse or promote products
    derived from this software without specific prior written permission.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
    EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
    DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
    (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
    LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
    ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#pragma once
#include <algorithm>
#include <memory>
#include <vector>
#include <map>
#include <string>
#include <limits>
#include <unordered_map>

#include "caffe.pb.h"

#include "tiny_dnn/layers/convolutional_layer.h"
#include "tiny_dnn/layers/deconvolutional_layer.h"
#include "tiny_dnn/layers/fully_connected_layer.h"
#include "tiny_dnn/layers/average_pooling_layer.h"
#include "tiny_dnn/layers/max_pooling_layer.h"
#include "tiny_dnn/layers/linear_layer.h"
#include "tiny_dnn/layers/lrn_layer.h"
#include "tiny_dnn/layers/dropout_layer.h"

typedef tiny_dnn::shape3d shape_t;

#ifdef _MSC_VER
#define _NOMINMAX
#include <io.h>
#include <fcntl.h>
#define CNN_OPEN_BINARY(filename) open(filename, _O_RDONLY|_O_BINARY)
#define CNN_OPEN_TXT(filename) open(filename, _O_RDONLY)
#pragma warning(push)
#pragma warning(disable:4996)
#else
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#define CNN_OPEN_BINARY(filename) open(filename, O_RDONLY)
#define CNN_OPEN_TXT(filename) open(filename, O_RDONLY)
#endif

namespace tiny_dnn {
namespace detail {

inline void read_proto_from_text(const std::string& prototxt,
                                 google::protobuf::Message *message) {
    int fd = CNN_OPEN_TXT(prototxt.c_str());
    if (fd == -1) {
        throw nn_error("file not fonud: " + prototxt);
    }

    google::protobuf::io::FileInputStream input(fd);
    input.SetCloseOnDelete(true);

    if (!google::protobuf::TextFormat::Parse(&input, message)) {
        throw nn_error("failed to parse");
    }
}

inline void read_proto_from_binary(const std::string& protobinary,
                                   google::protobuf::Message *message) {
    int fd = CNN_OPEN_BINARY(protobinary.c_str());
    google::protobuf::io::FileInputStream rawstr(fd);
    google::protobuf::io::CodedInputStream codedstr(&rawstr);

    rawstr.SetCloseOnDelete(true);
    codedstr.SetTotalBytesLimit(std::numeric_limits<int>::max(),
                                std::numeric_limits<int>::max() / 2);

    if (!message->ParseFromCodedStream(&codedstr)) {
        throw nn_error("failed to parse");
    }
}

inline std::shared_ptr<weight_init::function>
create_filler(const std::string& filler) {
    if (filler == "xavier") {
        return std::make_shared<weight_init::xavier>();
    } else if (filler == "constant") {
        return std::make_shared<weight_init::constant>();
    } else if (filler == "gaussian") {
        return std::make_shared<weight_init::gaussian>();
    } else {
        throw nn_error("unsupported filler type");
    }
}

template <typename param>
inline bool get_kernel_size_2d(const param& p, layer_size_t *kernel) {
    if (p.has_kernel_w() && p.has_kernel_h()) {
        if (p.kernel_w() != p.kernel_h()) {
            throw nn_error("unsupported kernel shape");
        }
        *kernel = p.kernel_w();
        return true;
    }
    return false;
}

template <typename param>
inline bool get_kernel_size_2d(const param& p, layer_size_t *kernel_w, layer_size_t *kernel_h) {
    if (p.has_kernel_w() && p.has_kernel_h()) {
        *kernel_w = p.kernel_w();
        *kernel_h = p.kernel_h();
        return true;
    }
    return false;
}

inline layer_size_t get_kernel_size_2d(const caffe::ConvolutionParameter& p) {
    layer_size_t window_size;
    if (!get_kernel_size_2d(p, &window_size)) {
        if (p.kernel_size_size() > 1) {
            throw nn_error("unsupported kernel shape");
        }
        window_size = p.kernel_size(0);
    }
    return window_size;
}

inline std::shared_ptr<layer> create_max_pool(layer_size_t pool_size_w,
                                              layer_size_t pool_size_h,
                                              layer_size_t stride_w,
                                              layer_size_t stride_h,
                                              padding pad_type,
                                              const shape_t& bottom_shape,
                                              shape_t *top_shape) {
    using max_pool = max_pooling_layer<activation::identity>;
    auto mp = std::make_shared<max_pool>(bottom_shape.width_,
                                         bottom_shape.height_,
                                         bottom_shape.depth_,
                                         pool_size_w, pool_size_h, stride_w, stride_h, pad_type);

    *top_shape = mp->out_shape()[0];
    mp->init_weight();

    return mp;
}

inline std::shared_ptr<layer> create_ave_pool(layer_size_t pool_size_w,
                                              layer_size_t pool_size_h,
                                              layer_size_t stride_w,
                                              layer_size_t stride_h,
                                              padding pad_type,
                                              const shape_t& bottom_shape,
                                              shape_t *top_shape) {
    using ave_pool = average_pooling_layer<activation::identity>;
    auto ap = std::make_shared<ave_pool>(bottom_shape.width_,
                                         bottom_shape.height_,
                                         bottom_shape.depth_,
                                         pool_size_w, pool_size_h, stride_w, stride_h, pad_type);

    // tiny-dnn has trainable parameter in average-pooling layer
    float_t weight = float_t(1) / (pool_size_w * pool_size_h);

    vec_t& w = *ap->weights()[0];
    vec_t& b = *ap->weights()[1];

    std::fill(w.begin(), w.end(), weight);
    std::fill(b.begin(), b.end(), float_t(0));

    *top_shape = ap->out_shape()[0];
    ap->init_weight();
    ap->set_trainable(false);

    return ap;
}

inline
std::shared_ptr<layer> create_softmax(const caffe::LayerParameter& layer,
                                      const shape_t& bottom_shape, shape_t *) {
    auto sm = std::make_shared<linear_layer<activation::softmax>>(
        bottom_shape.size());
    sm->init_weight();
    return sm;
}

inline
std::shared_ptr<layer> create_sigmoid(const caffe::LayerParameter& layer,
                                      const shape_t& bottom_shape, shape_t *) {
    auto ce = std::make_shared<linear_layer<activation::sigmoid>>(
        bottom_shape.size());
    return ce;
}

inline
std::shared_ptr<layer> create_tanh(const caffe::LayerParameter& layer,
                                   const shape_t& bottom_shape, shape_t *) {
    auto tanh = std::make_shared<linear_layer<activation::tan_h>>(
        bottom_shape.size());
    return tanh;
}

inline
std::shared_ptr<layer> create_power(const caffe::LayerParameter& layer,
                                    const shape_t& bottom_shape, shape_t *) {
    auto power = std::make_shared<power_layer>(bottom_shape, layer.power_param().power(), layer.power_param().scale());
    return power;
}

    
inline
std::shared_ptr<layer> create_pooling(const caffe::LayerParameter& layer,
                                      const shape_t& bottom_shape,
                                      shape_t *top_shape) {
    if (!layer.has_pooling_param()) {
        throw nn_error("pool param missing");
    }

    auto pool_param = layer.pooling_param();

    layer_size_t h_stride = 0;
    layer_size_t w_stride = 0;
    layer_size_t pool_size_w = 0;
    layer_size_t pool_size_h = 0;
    layer_size_t h_pad = 0;
    layer_size_t w_pad = 0;
    padding pad_type = padding::valid;

    if (!get_kernel_size_2d(pool_param, &pool_size_w, &pool_size_h)) {
        pool_size_w = pool_size_h = pool_param.kernel_size();
    }

    if (pool_param.has_stride() || pool_param.has_stride_h()) {
        h_stride = pool_param.has_stride() ?
                   pool_param.stride() : pool_param.stride_h();
    }

    if (pool_param.has_stride() || pool_param.has_stride_w()) {
        w_stride = pool_param.has_stride() ?
                   pool_param.stride() : pool_param.stride_w();
    }

    if (pool_param.has_pad() || pool_param.has_pad_w()) {
        w_pad = pool_param.has_pad() ?
                pool_param.pad() : pool_param.pad_w();
    }

    if (pool_param.has_pad() || pool_param.has_pad_h()) {
        h_pad = pool_param.has_pad() ?
                pool_param.pad() : pool_param.pad_h();
    }

    if (w_pad != 0) {
        if (w_pad == pool_size_w - 1) {
            pad_type = padding::same;
        }
        else {
            throw nn_error("unsupported padding type");
        }
    }

    if (h_pad != 0) {
        if (h_pad == pool_size_h - 1) {
            pad_type = padding::same;
        }
        else {
            throw nn_error("unsupported padding type");
        }
    }

    if (pool_param.has_pool()) {
        auto type = pool_param.pool();

        switch (type) {
            case caffe::PoolingParameter_PoolMethod_MAX:
                return create_max_pool(pool_size_w, pool_size_h, w_stride, h_stride, pad_type,
                                       bottom_shape, top_shape);
            case caffe::PoolingParameter_PoolMethod_AVE:
                return create_ave_pool(pool_size_w, pool_size_h, w_stride, h_stride, pad_type,
                                       bottom_shape, top_shape);
            default:
                throw nn_error("unsupported layer type");
        }
    }

    // default: max-pool
    return create_max_pool(pool_size_w, pool_size_h, w_stride, h_stride, pad_type, bottom_shape, top_shape);
}

inline
std::shared_ptr<layer> create_relu(const caffe::LayerParameter& layer,
                                   const shape_t& bottom_shape, shape_t *) {
    auto relu = std::make_shared<linear_layer<activation::relu>>(
        bottom_shape.size());
    return relu;
}

inline std::shared_ptr<layer> create_batchnorm(const caffe::LayerParameter& layer,
    const shape_t& bottom_shape, shape_t *top_shape) {
    using bn_layer = batch_normalization_layer;

    *top_shape = bottom_shape;

    float_t eps = 1e-5f;
    float_t momentum = 0.999f;

    if (layer.has_batch_norm_param()) {
        auto bn_param = layer.batch_norm_param();

        if (bn_param.has_eps()) {
            eps = bn_param.eps();
        }
        if (bn_param.has_moving_average_fraction()) {
            momentum = bn_param.moving_average_fraction();
        }
    }

    auto bn = std::make_shared<bn_layer>(bottom_shape.area(), bottom_shape.depth_, eps, momentum, net_phase::test);

    // weight
    if (layer.blobs_size() > 0) {
        auto global_stats = layer.blobs();
        if (global_stats.size() != 3) {
            throw std::runtime_error("unexpected bn stored statistics");
        }       

        float_t scale_factor = global_stats.Get(2).data(0) == 0 ? 0 : 1 / global_stats.Get(2).data(0);
        vec_t mean(bottom_shape.depth_);
        vec_t variance(bottom_shape.depth_);

        for (size_t i = 0; i < mean.size(); i++) {
            mean[i]     = global_stats.Get(0).data(i) * scale_factor;
            variance[i] = global_stats.Get(1).data(i) * scale_factor;
        }
        bn->set_mean(mean);
        bn->set_variance(variance);
    }

    return bn;
}


inline void load_weights_fullyconnected(const caffe::LayerParameter& src,
                                        layer *dst) {
    auto weights = src.blobs(0);
    int curr = 0;

    const auto dst_out_size = dst->out_size();
    const auto dst_in_size = dst->in_size();

    if (dst_out_size * dst_in_size !=
        static_cast<serial_size_t>(weights.data_size())) {
        throw nn_error(
            std::string("layer size mismatch!") +
            "caffe(" + src.name() + "):" + to_string(weights.data_size()) + "\n" +
            "tiny-dnn(" + dst->layer_type() + "):" + to_string(dst->weights().size()));
    }

    vec_t& w = *dst->weights()[0];
    vec_t& b = *dst->weights()[1];

    // fill weights
    for (size_t o = 0; o < dst_out_size; o++) {
        for (size_t i = 0; i < dst_in_size; i++) {
            // TODO: how to access to weights?
            //dst->weight()[i * dst->out_size() + o] = weights.data(curr++); // transpose
            w[i * dst_out_size + o] = weights.data(curr++); // transpose
        }
    }

    // fill bias
    if (src.inner_product_param().bias_term()) {
        auto biases = src.blobs(1);
        for (size_t o = 0; o < dst_out_size; o++) {
            // TODO: how to access to biases?
            //dst->bias()[o] = biases.data(o);
            b[o] = biases.data(o);
        }
    }
}

inline std::shared_ptr<layer> create_fullyconnected(
        const caffe::LayerParameter& layer,
        const shape_t& bottom_shape, shape_t *top_shape) {
    using fc_layer = fully_connected_layer<activation::identity>;

    if (!layer.has_inner_product_param()) {
        throw nn_error("inner-product param missing");
    }

    layer_size_t dim_input = 0, dim_output = 0;
    bool has_bias = true;

    auto ip_param = layer.inner_product_param();
    has_bias = ip_param.bias_term();

    dim_output = ip_param.num_output();
    dim_input = bottom_shape.size();

    auto ip = std::make_shared<fc_layer>(dim_input, dim_output, has_bias);

    // filler
    if (ip_param.has_weight_filler()) {
        ip->weight_init(create_filler(ip_param.weight_filler().type()));
    }

    if (ip_param.has_bias_filler()) {
        ip->bias_init(create_filler(ip_param.bias_filler().type()));
    }

    // weight
    if (layer.blobs_size() > 0) {
        load_weights_fullyconnected(layer, ip.get());
    }

    // TODO: check if it works
    *top_shape = ip->out_shape()[0];
    return ip;
}

inline void load_weights_conv(const caffe::LayerParameter& src, layer *dst) {
    // fill weight
    auto weights = src.blobs(0);

    //TODO: check if it works
    //int out_channels = dst->out_shape().depth_;
    //int in_channels = dst->in_shape().depth_;
    int out_channels = dst->out_data_shape()[0].depth_;
    int in_channels = dst->in_data_shape()[0].depth_;

    connection_table table;
    auto conv_param = src.convolution_param();
    int dst_idx = 0;
    int src_idx = 0;
    int window_size = get_kernel_size_2d(conv_param);

    if (conv_param.has_group()) {
        table = connection_table(conv_param.group(), in_channels, out_channels);
    }

    vec_t& w = *dst->weights()[0];
    vec_t& b = *dst->weights()[1];

    // fill weights
    for (int o = 0; o < out_channels; o++) {
        for (int i = 0; i < in_channels; i++) {
            if (!table.is_connected(o, i)) {
                dst_idx += window_size * window_size;
                continue;
            }
            for (int x = 0; x < window_size * window_size; x++) {
                //TODO
                //dst->weight()[dst_idx++] = weights.data(src_idx++);
                w[dst_idx++] =  weights.data(src_idx++);
            }
        }
    }

    // fill bias
    if (conv_param.bias_term()) {
        auto biases = src.blobs(1);
        for (int o = 0; o < out_channels; o++) {
            //TODO
            //dst->bias()[o] = biases.data(o);
            b[o] = biases.data(o);
        }
    }
}

inline void load_weights_pool(const caffe::LayerParameter& src, layer *dst) {
    auto pool_param = src.pooling_param();

    //TODO
    //if (dst->weight().size()) {
    if (dst->weights().size()) {
        layer_size_t pool_size = 0;

        if (!get_kernel_size_2d(pool_param, &pool_size)) {
            pool_size = pool_param.kernel_size();
        }

        // tiny-dnn has trainable parameter in average-pooling layer
        float_t weight = float_t(1) / sqr(pool_size);

        //TODO
        /*if (!dst->weight().empty()) {
            std::fill(dst->weight().begin(), dst->weight().end(), weight);
        }
        if (!dst->bias().empty()) {
            std::fill(dst->bias().begin(), dst->bias().end(), float_t(0));
            dst->init_bias();
        }*/

        vec_t& w = *dst->weights()[0];
        vec_t& b = *dst->weights()[1];

        if (!w.empty()) {
            std::fill(w.begin(), w.end(), weight);
        }
        if (!b.empty()) {
            std::fill(b.begin(), b.end(), float_t(0));
            //dst->init_bias();
        }
    }
}

inline
std::shared_ptr<layer> create_lrn(const caffe::LayerParameter& layer,
                                  const shape_t& bottom_shape,
                                  shape_t *top_shape) {
    using lrn_layer = lrn_layer<activation::identity>;

    if (!layer.has_lrn_param()) {
        throw nn_error("lrn param missing");
    }

    auto lrn_param = layer.lrn_param();
    layer_size_t local_size = 5;
    float_t alpha = 1;
    float_t beta = 5;
    norm_region region = norm_region::across_channels;

    if (lrn_param.has_local_size()) local_size = lrn_param.local_size();
    if (lrn_param.has_alpha()) alpha = lrn_param.alpha();
    if (lrn_param.has_beta()) beta = lrn_param.beta();
    if (lrn_param.has_norm_region()) {
        if (lrn_param.norm_region() == caffe::LRNParameter_NormRegion_WITHIN_CHANNEL) // NOLINT
            region = norm_region::within_channels;
    }

    auto lrn = std::make_shared<lrn_layer>(bottom_shape.width_,
                                           bottom_shape.height_,
                                           local_size,
                                           bottom_shape.depth_,
                                           alpha, beta, region);
    return lrn;
}

inline
std::shared_ptr<layer> create_dropout(const caffe::LayerParameter& layer,
                                      const shape_t& bottom_shape,
                                      shape_t *top_shape) {
    if (!layer.has_dropout_param()) {
        throw nn_error("dropout param missing");
    }

    float_t dropout_rate = float_t(0.5);

    if (layer.dropout_param().has_dropout_ratio()) {
        dropout_rate = layer.dropout_param().dropout_ratio();
    }

    auto dropout = std::make_shared<dropout_layer>(bottom_shape.size(),
                                                   dropout_rate,
                                                   net_phase::test);
    return dropout;
}

inline
std::shared_ptr<layer> create_convlayer(const caffe::LayerParameter& layer,
                                        const shape_t& bottom_shape,
                                        shape_t *top_shape) {
    using conv_layer = convolutional_layer<activation::identity>;

    if (!layer.has_convolution_param()) {
        throw nn_error("convolution param missing");
    }

    // layer parameters
    layer_size_t in_width = 0, in_height = 0, window_size = 0;
    layer_size_t in_channels = 0, out_channels = 0;
    layer_size_t w_stride = 1, h_stride = 1;
    bool has_bias = true;
    padding pad_type = padding::valid;
    connection_table table;

    auto conv_param = layer.convolution_param();

    // shape
    out_channels = conv_param.num_output();
    in_channels = bottom_shape.depth_;
    in_width = bottom_shape.width_;
    in_height = bottom_shape.height_;
    has_bias = conv_param.bias_term();
    window_size = get_kernel_size_2d(conv_param);

    // padding
    if (conv_param.pad_size() == 1 ||
       (conv_param.has_pad_w() && conv_param.has_pad_h())) {
        uint32_t pad_w = conv_param.pad_size() == 1 ?
                         conv_param.pad(0) : conv_param.pad_w();

        uint32_t pad_h = conv_param.pad_size() == 1 ?
                         conv_param.pad(0) : conv_param.pad_h();

        if (pad_w != pad_h) {
            throw nn_error("conv:not supported padding size");
        }

        // 0 ... valid, (window_size-1)/2 ... same
        if (pad_w == (window_size - 1) / 2) {
            pad_type = padding::same;
        } else if (pad_w == 0) {
            pad_type = padding::valid;
        } else {
            throw nn_error("conv:not supported padding size");
        }
    }

    // stride
    if (conv_param.stride_size() == 1 || conv_param.has_stride_h()) {
        h_stride = conv_param.stride_size() == 1 ?
                   conv_param.stride(0) : conv_param.stride_h();
    }

    if (conv_param.stride_size() == 1 || conv_param.has_stride_w()) {
        w_stride = conv_param.stride_size() == 1 ?
                   conv_param.stride(0) : conv_param.stride_w();
    }

    // group
    if (conv_param.has_group()) {
        table = connection_table(conv_param.group(), in_channels, out_channels);
    }

    auto conv = std::make_shared<conv_layer>(in_width, in_height,
                                             window_size,
                                             in_channels, out_channels,
                                             table,
                                             pad_type,
                                             has_bias,
                                             w_stride, h_stride);
    // filler
    if (conv_param.has_weight_filler()) {
        conv->weight_init(create_filler(conv_param.weight_filler().type()));
    }

    if (conv_param.has_bias_filler()) {
        conv->bias_init(create_filler(conv_param.bias_filler().type()));
    }

    // set weight (optional)
    if (layer.blobs_size() > 0) {  // blobs(0)...weight, blobs(1)...bias
        load_weights_conv(layer, conv.get());
    }

    *top_shape = conv->out_shape()[0];
    return conv;
}

inline
std::shared_ptr<layer> create_deconvlayer(const caffe::LayerParameter& layer,
                                        const shape_t& bottom_shape,
                                        shape_t *top_shape) {
    using deconv_layer = deconvolutional_layer<activation::identity>;

    if (!layer.has_convolution_param()) {
        throw nn_error("deconvolution param missing");
    }

    // layer parameters
    layer_size_t in_width = 0, in_height = 0, window_size = 0;
    layer_size_t in_channels = 0, out_channels = 0;
    layer_size_t w_stride = 1, h_stride = 1;
    bool has_bias = true;
    padding pad_type = padding::valid;
    connection_table table;

    auto deconv_param = layer.convolution_param();

    // shape
    out_channels = deconv_param.num_output();
    in_channels = bottom_shape.depth_;
    in_width = bottom_shape.width_;
    in_height = bottom_shape.height_;
    has_bias = deconv_param.bias_term();
    window_size = get_kernel_size_2d(deconv_param);

    // unpadding
    if (deconv_param.pad_size() == 1 ||
       (deconv_param.has_pad_w() && deconv_param.has_pad_h())) {
        uint32_t unpad_w = deconv_param.pad_size() == 1 ?
                         deconv_param.pad(0) : deconv_param.pad_w();

        uint32_t unpad_h = deconv_param.pad_size() == 1 ?
                         deconv_param.pad(0) : deconv_param.pad_h();

        if (unpad_w != unpad_h) {
            throw nn_error("deconv:not supported unpadding size");
        }

        // 0 ... valid, (window_size-1)/2 ... same
        if (unpad_w == (window_size - 1) / 2) {
            pad_type = padding::same;
        } else if (unpad_w == 0) {
            pad_type = padding::valid;
        } else {
            throw nn_error("deconv:not supported unpadding size");
        }
    }

    // stride
    if (deconv_param.stride_size() == 1 || deconv_param.has_stride_h()) {
        h_stride = deconv_param.stride_size() == 1 ?
                   deconv_param.stride(0) : deconv_param.stride_h();
    }

    if (deconv_param.stride_size() == 1 || deconv_param.has_stride_w()) {
        w_stride = deconv_param.stride_size() == 1 ?
                   deconv_param.stride(0) : deconv_param.stride_w();
    }

    // group
    if (deconv_param.has_group()) {
        table = connection_table(deconv_param.group(), in_channels, out_channels);
    }

    auto deconv = std::make_shared<deconv_layer>(in_width, in_height,
                                             window_size,
                                             in_channels, out_channels,
                                             table,
                                             pad_type,
                                             has_bias,
                                             w_stride, h_stride);
    // filler
    if (deconv_param.has_weight_filler()) {
        deconv->weight_init(create_filler(deconv_param.weight_filler().type()));
    }

    if (deconv_param.has_bias_filler()) {
        deconv->bias_init(create_filler(deconv_param.bias_filler().type()));
    }

    // set weight (optional)
    if (layer.blobs_size() > 0) {  // blobs(0)...weight, blobs(1)...bias
        load_weights_conv(layer, deconv.get());
    }
    //TODO
    //*top_shape = deconv->out_shape();
    *top_shape = deconv->out_shape()[0];
    return deconv;
}

inline bool layer_skipped(const std::string& type) {
    if (type == "Data" || type == "EuclideanLoss" || type == "Input") return true;
    return false;
}

inline bool layer_supported(const std::string& type) {
    static const char* supported[] = {
        "InnerProduct", "Convolution", "Deconvolution", "Pooling",
        "LRN", "Dropout",
        "SoftmaxWithLoss", "SigmoidCrossEntropyLoss",
        "ReLU", "Sigmoid", "TanH", "Softmax", "BatchNorm", "Power"
    };

    for (size_t i = 0; i < sizeof(supported) / sizeof(supported[0]); i++) {
        if (supported[i] == type) return true;
    }
    return false;
}

inline bool layer_match(const std::string& caffetype,
                        const std::string& tiny_dnn_type) {
    const char* conversions[][2] = {
        { "InnerProduct", "fully-connected" },
        { "Convolution", "conv" },
        { "Deconvolution", "deconv" },
        { "Pooling", "ave-pool" },
        { "Pooling", "max-pool" }
    };

    for (size_t i = 0; i < sizeof(conversions) / sizeof(conversions[0]); i++) {
        if (conversions[i][0] == caffetype &&
            conversions[i][1] == tiny_dnn_type) return true;
    }
    return false;
}

inline std::shared_ptr<layer> create(const caffe::LayerParameter& layer,
                                     const shape_t& in_shape,
                                     shape_t* out_shape) {
    const std::string layer_type = layer.type();

    if (layer_type == "Convolution") {
        return detail::create_convlayer(layer, in_shape, out_shape);
    }

    if (layer_type == "Deconvolution") {
        return detail::create_deconvlayer(layer, in_shape, out_shape);
    }

    if (layer_type == "InnerProduct") {
        return detail::create_fullyconnected(layer, in_shape, out_shape);
    }

    if (layer_type == "Pooling") {
        return detail::create_pooling(layer, in_shape, out_shape);
    }

    if (layer_type == "BatchNorm") {
        return detail::create_batchnorm(layer, in_shape, out_shape);
    }

    if (layer_type == "LRN") {
        return detail::create_lrn(layer, in_shape, out_shape);
    }

    if (layer_type == "Dropout") {
        return detail::create_dropout(layer, in_shape, out_shape);
    }

    if (layer_type == "SoftmaxWithLoss" ||
        layer_type == "Softmax") {
        return detail::create_softmax(layer, in_shape, out_shape);
    }

    if (layer_type == "SigmoidCrossEntropyLoss" ||
        layer_type == "Sigmoid") {
        return detail::create_sigmoid(layer, in_shape, out_shape);
    }

    if (layer_type == "ReLU") {
        return detail::create_relu(layer, in_shape, out_shape);
    }

    if (layer_type == "TanH") {
        return detail::create_tanh(layer, in_shape, out_shape);
    }
    
    if (layer_type == "Power") {
        return detail::create_power(layer, in_shape, out_shape);
    }

    throw nn_error("layer parser not found");
}

inline void load(const caffe::LayerParameter& src, layer *dst) {
    typedef std::function<void(const caffe::LayerParameter&, layer*)> factoryimpl; // NOLINT
    std::unordered_map<std::string, factoryimpl> factory_registry;

    factory_registry["Convolution"] = detail::load_weights_conv;
    factory_registry["Deconvolution"] = detail::load_weights_conv;
    factory_registry["InnerProduct"] = detail::load_weights_fullyconnected;
    factory_registry["Pooling"] = detail::load_weights_pool;

    if (factory_registry.find(src.type()) == factory_registry.end()) {
        throw nn_error("layer parser not found");
    }

    return factory_registry[src.type()](src, dst);
}


struct layer_node {
    const caffe::LayerParameter *layer;
    const layer_node *next;  // top-side
    const layer_node *prev;  // bottom-side

    layer_node() : layer(0), next(0), prev(0) {}
    explicit layer_node(const caffe::LayerParameter *l)
        : layer(l), next(0), prev(0) {}
};

// parse caffe net and interpret as single layer vector
class caffe_layer_vector {
 public:
    explicit caffe_layer_vector(const caffe::NetParameter& net_orig)
            : net(net_orig) {
        if (net.layers_size() > 0) {
            upgradev1net(net_orig, &net);
        }

        nodes.reserve(net.layer_size());

        for (int i = 0; i < net.layer_size(); i++) {
            auto& l = net.layer(i);

            if (layer_table.find(l.name()) != layer_table.end()) continue;

            nodes.emplace_back(&l);
            layer_table[l.name()] = &nodes.back();
        }

        for (size_t i = 0; i < nodes.size(); i++) {
            auto& l = nodes[i];

            if (l.layer->bottom_size() > 0 && blob_table[l.layer->bottom(0)]) {
                auto& bottom = blob_table[l.layer->bottom(0)];
                l.prev = bottom;
                layer_table[bottom->layer->name()]->next = &l;
            }

            if (l.layer->top_size() > 0) {
                blob_table[l.layer->top(0)] = &l;
            }
        }

        auto root = std::find_if(nodes.begin(),
                                 nodes.end(), [](const layer_node& n) {
            return n.prev == 0;
        });

        if (root == nodes.end()) {
            throw nn_error("root layer not found");
        }

        root_node = &*root;
        const layer_node *current = &*root;

        while (current) {
            node_list.push_back(current->layer);
            current = current->next;
        }
    }

    size_t size() const {
        return node_list.size();
    }

    const caffe::LayerParameter& operator[] (size_t index) const {
        return *(node_list[index]);
    }

 private:
    void upgradev1net(const caffe::NetParameter& old,
                      caffe::NetParameter *dst) const {
        dst->CopyFrom(old);
        dst->clear_layers();
        dst->clear_layer();

        for (int i = 0; i < old.layers_size(); i++) {
            upgradev1layer(old.layers(i), dst->add_layer());
        }
    }

    const char* v1type2name(caffe::V1LayerParameter_LayerType type) const {
        switch (type) {
        case caffe::V1LayerParameter_LayerType_NONE:
            return "";
        case caffe::V1LayerParameter_LayerType_ABSVAL:
            return "AbsVal";
        case caffe::V1LayerParameter_LayerType_ACCURACY:
            return "Accuracy";
        case caffe::V1LayerParameter_LayerType_ARGMAX:
            return "ArgMax";
        case caffe::V1LayerParameter_LayerType_BNLL:
            return "BNLL";
        case caffe::V1LayerParameter_LayerType_CONCAT:
            return "Concat";
        case caffe::V1LayerParameter_LayerType_CONTRASTIVE_LOSS:
            return "ContrastiveLoss";
        case caffe::V1LayerParameter_LayerType_CONVOLUTION:
            return "Convolution";
        case caffe::V1LayerParameter_LayerType_DECONVOLUTION:
            return "Deconvolution";
        case caffe::V1LayerParameter_LayerType_DATA:
            return "Data";
        case caffe::V1LayerParameter_LayerType_DROPOUT:
            return "Dropout";
        case caffe::V1LayerParameter_LayerType_DUMMY_DATA:
            return "DummyData";
        case caffe::V1LayerParameter_LayerType_EUCLIDEAN_LOSS:
            return "EuclideanLoss";
        case caffe::V1LayerParameter_LayerType_ELTWISE:
            return "Eltwise";
        case caffe::V1LayerParameter_LayerType_EXP:
            return "Exp";
        case caffe::V1LayerParameter_LayerType_FLATTEN:
            return "Flatten";
        case caffe::V1LayerParameter_LayerType_HDF5_DATA:
            return "HDF5Data";
        case caffe::V1LayerParameter_LayerType_HDF5_OUTPUT:
            return "HDF5Output";
        case caffe::V1LayerParameter_LayerType_HINGE_LOSS:
            return "HingeLoss";
        case caffe::V1LayerParameter_LayerType_IM2COL:
            return "Im2col";
        case caffe::V1LayerParameter_LayerType_IMAGE_DATA:
            return "ImageData";
        case caffe::V1LayerParameter_LayerType_INFOGAIN_LOSS:
            return "InfogainLoss";
        case caffe::V1LayerParameter_LayerType_INNER_PRODUCT:
            return "InnerProduct";
        case caffe::V1LayerParameter_LayerType_LRN:
            return "LRN";
        case caffe::V1LayerParameter_LayerType_MEMORY_DATA:
            return "MemoryData";
        case caffe::V1LayerParameter_LayerType_MULTINOMIAL_LOGISTIC_LOSS:
            return "MultinomialLogisticLoss";
        case caffe::V1LayerParameter_LayerType_MVN:
            return "MVN";
        case caffe::V1LayerParameter_LayerType_POOLING:
            return "Pooling";
        case caffe::V1LayerParameter_LayerType_POWER:
            return "Power";
        case caffe::V1LayerParameter_LayerType_RELU:
            return "ReLU";
        case caffe::V1LayerParameter_LayerType_SIGMOID:
            return "Sigmoid";
        case caffe::V1LayerParameter_LayerType_SIGMOID_CROSS_ENTROPY_LOSS:
            return "SigmoidCrossEntropyLoss";
        case caffe::V1LayerParameter_LayerType_SILENCE:
            return "Silence";
        case caffe::V1LayerParameter_LayerType_SOFTMAX:
            return "Softmax";
        case caffe::V1LayerParameter_LayerType_SOFTMAX_LOSS:
            return "SoftmaxWithLoss";
        case caffe::V1LayerParameter_LayerType_SPLIT:
            return "Split";
        case caffe::V1LayerParameter_LayerType_SLICE:
            return "Slice";
        case caffe::V1LayerParameter_LayerType_TANH:
            return "TanH";
        case caffe::V1LayerParameter_LayerType_WINDOW_DATA:
            return "WindowData";
        case caffe::V1LayerParameter_LayerType_THRESHOLD:
            return "Threshold";
        default:
            throw nn_error("unknown v1 layer-type");
        }
    }

    void upgradev1layer(const caffe::V1LayerParameter& old,
                        caffe::LayerParameter *dst) const {
        dst->Clear();

        for (int i = 0; i < old.bottom_size(); i++) {
            dst->add_bottom(old.bottom(i));
        }

        for (int i = 0; i < old.top_size(); i++) {
            dst->add_top(old.top(i));
        }

        if (old.has_name()) dst->set_name(old.name());
        if (old.has_type()) dst->set_type(v1type2name(old.type()));

        for (int i = 0; i < old.blobs_size(); i++) {
            dst->add_blobs()->CopyFrom(old.blobs(i));
        }

        for (int i = 0; i < old.param_size(); i++) {
            while (dst->param_size() <= i) dst->add_param();
            dst->mutable_param(i)->set_name(old.param(i));
        }

        #define COPY_PARAM(name) if (old.has_##name##_param()) dst->mutable_##name##_param()->CopyFrom(old.name##_param())

        COPY_PARAM(accuracy);
        COPY_PARAM(argmax);
        COPY_PARAM(concat);
        COPY_PARAM(contrastive_loss);
        COPY_PARAM(convolution);
        COPY_PARAM(data);
        COPY_PARAM(dropout);
        COPY_PARAM(dummy_data);
        COPY_PARAM(eltwise);
        COPY_PARAM(exp);
        COPY_PARAM(hdf5_data);
        COPY_PARAM(hdf5_output);
        COPY_PARAM(hinge_loss);
        COPY_PARAM(image_data);
        COPY_PARAM(infogain_loss);
        COPY_PARAM(inner_product);
        COPY_PARAM(lrn);
        COPY_PARAM(memory_data);
        COPY_PARAM(mvn);
        COPY_PARAM(pooling);
        COPY_PARAM(power);
        COPY_PARAM(relu);
        COPY_PARAM(sigmoid);
        COPY_PARAM(softmax);
        COPY_PARAM(slice);
        COPY_PARAM(tanh);
        COPY_PARAM(threshold);
        COPY_PARAM(window_data);
        COPY_PARAM(transform);
        COPY_PARAM(loss);
        #undef COPY_PARAM
    }

    caffe::NetParameter net;
    layer_node *root_node;
    /* layer name -> layer */
    std::map<std::string, layer_node*> layer_table;
    /* blob name -> bottom holder */
    std::map<std::string, layer_node*> blob_table;
    std::vector<layer_node> nodes;
    std::vector<const caffe::LayerParameter*> node_list;
};

}  // namespace detail
}  // namespace tiny_dnn

#ifdef _MSC_VER
#pragma warning(pop)
#endif
