#include <string>

#include "caffe/layer.hpp"
#include "caffe/layer_factory.hpp"
#include "caffe/proto/caffe.pb.h"
#include "caffe/vision_layers.hpp"

#ifdef WITH_PYTHON_LAYER
#include "caffe/python_layer.hpp"
#endif

namespace caffe {

// Get convolution layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetConvolutionLayer(
    const LayerParameter& param) {
  ConvolutionParameter_Engine engine = param.convolution_param().engine();
  if (engine == ConvolutionParameter_Engine_DEFAULT) {
    engine = ConvolutionParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = ConvolutionParameter_Engine_CUDNN;
#endif
  }
  if (engine == ConvolutionParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new ConvolutionLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == ConvolutionParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNConvolutionLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
  }
    return nullptr;
}

REGISTER_LAYER_CREATOR(Convolution, GetConvolutionLayer);

// Get pooling layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
  PoolingParameter_Engine engine = param.pooling_param().engine();
  if (engine == PoolingParameter_Engine_DEFAULT) {
    engine = PoolingParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = PoolingParameter_Engine_CUDNN;
#endif
  }
  if (engine == PoolingParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == PoolingParameter_Engine_CUDNN) {
    PoolingParameter p_param = param.pooling_param();
    if (p_param.pad() || p_param.pad_h() || p_param.pad_w() ||
        param.top_size() > 1) {
      LOG(INFO) << "CUDNN does not support padding or multiple tops. "
                << "Using Caffe's own pooling layer.";
      return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
    }
    return shared_ptr<Layer<Dtype> >(new CuDNNPoolingLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
  }
        return nullptr;
}

REGISTER_LAYER_CREATOR(Pooling, GetPoolingLayer);

// Get relu layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetReLULayer(const LayerParameter& param) {
  ReLUParameter_Engine engine = param.relu_param().engine();
  if (engine == ReLUParameter_Engine_DEFAULT) {
    engine = ReLUParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = ReLUParameter_Engine_CUDNN;
#endif
  }
  if (engine == ReLUParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new ReLULayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == ReLUParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNReLULayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
  }
        return nullptr;
}

REGISTER_LAYER_CREATOR(ReLU, GetReLULayer);

    
//#ifndef COMPACT
//// Get sigmoid layer according to engine.
//template <typename Dtype>
//shared_ptr<Layer<Dtype> > GetSigmoidLayer(const LayerParameter& param) {
//  SigmoidParameter_Engine engine = param.sigmoid_param().engine();
//  if (engine == SigmoidParameter_Engine_DEFAULT) {
//    engine = SigmoidParameter_Engine_CAFFE;
//#ifdef USE_CUDNN
//    engine = SigmoidParameter_Engine_CUDNN;
//#endif
//  }
//  if (engine == SigmoidParameter_Engine_CAFFE) {
//    return shared_ptr<Layer<Dtype> >(new SigmoidLayer<Dtype>(param));
//#ifdef USE_CUDNN
//  } else if (engine == SigmoidParameter_Engine_CUDNN) {
//    return shared_ptr<Layer<Dtype> >(new CuDNNSigmoidLayer<Dtype>(param));
//#endif
//  } else {
//    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
//  }
//        return nullptr;
//}
//
//REGISTER_LAYER_CREATOR(Sigmoid, GetSigmoidLayer);
//#endif
    
// Get softmax layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetSoftmaxLayer(const LayerParameter& param) {
  SoftmaxParameter_Engine engine = param.softmax_param().engine();
  if (engine == SoftmaxParameter_Engine_DEFAULT) {
    engine = SoftmaxParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = SoftmaxParameter_Engine_CUDNN;
#endif
  }
  if (engine == SoftmaxParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new SoftmaxLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == SoftmaxParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNSoftmaxLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
  }
        return nullptr;
}

REGISTER_LAYER_CREATOR(Softmax, GetSoftmaxLayer);

//#ifndef COMPACT
//// Get tanh layer according to engine.
//template <typename Dtype>
//shared_ptr<Layer<Dtype> > GetTanHLayer(const LayerParameter& param) {
//  TanHParameter_Engine engine = param.tanh_param().engine();
//  if (engine == TanHParameter_Engine_DEFAULT) {
//    engine = TanHParameter_Engine_CAFFE;
//#ifdef USE_CUDNN
//    engine = TanHParameter_Engine_CUDNN;
//#endif
//  }
//  if (engine == TanHParameter_Engine_CAFFE) {
//    return shared_ptr<Layer<Dtype> >(new TanHLayer<Dtype>(param));
//#ifdef USE_CUDNN
//  } else if (engine == TanHParameter_Engine_CUDNN) {
//    return shared_ptr<Layer<Dtype> >(new CuDNNTanHLayer<Dtype>(param));
//#endif
//  } else {
//    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
//  }
//        return nullptr;
//}
//
//REGISTER_LAYER_CREATOR(TanH, GetTanHLayer);
//#endif
    
#ifdef WITH_PYTHON_LAYER
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetPythonLayer(const LayerParameter& param) {
  Py_Initialize();
  try {
    bp::object module = bp::import(param.python_param().module().c_str());
    bp::object layer = module.attr(param.python_param().layer().c_str())(param);
    return bp::extract<shared_ptr<PythonLayer<Dtype> > >(layer)();
  } catch (bp::error_already_set) {
    PyErr_Print();
    throw;
  }
}

REGISTER_LAYER_CREATOR(Python, GetPythonLayer);
#endif

// Layers that use their constructor as their default creator should be
// registered in their corresponding cpp files. Do not register them here.
}  // namespace caffe