#include <vector>

#include "caffe/layer.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/vision_layers.hpp"

namespace caffe {

template <typename Dtype>
void DeConcatLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  deconcat_dim_ = this->layer_param_.deconcat_param().deconcat_dim();
  CHECK_EQ(deconcat_dim_, 0) <<
    "deconcat_dim should be 0 for now";
}

template <typename Dtype>
void DeConcatLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  // Initialize with the first blob.
	 channels_ = bottom[0]->channels();
	  height_ = bottom[0]->height();
	  width_ = bottom[0]->width();
	  if (bottom.size()==1) {
  count_ = bottom[0]->count()/(top->size());
  num_ = bottom[0]->num()/(top->size());
  for (int i=0;i<top->size();i++) {
	  (*top)[i]->Reshape(num_, channels_, height_, width_);
	  CHECK_EQ(count_, (*top)[i]->count());
  }
	  } else {
		  const Dtype* bottom_data1=bottom[1]->cpu_data();
		  CHECK_EQ(bottom[1]->num(),top->size());
		  for (int i=0;i<top->size();i++) {
		  	  (*top)[i]->Reshape(int(bottom_data1[i]), channels_, height_, width_);
		    }
	  }
}

template <typename Dtype>
void DeConcatLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  const Dtype* bottom_data = bottom[0]->cpu_data();
  if (deconcat_dim_== 0) {
    int offset_num = 0;
    for (int i=0;i<top->size(); ++i) {
    	Dtype* top_data = (*top)[i]->mutable_cpu_data();
    	int num_elem = (*top)[i]->count();
    	caffe_copy(num_elem, bottom_data+bottom[0]->offset(offset_num), top_data);
    	offset_num += (*top)[i]->num();
    }
  }
}

template <typename Dtype>
void DeConcatLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
  //const Dtype* top_diff = top[0]->cpu_diff();
  Dtype* bottom_diff = (*bottom)[0]->mutable_cpu_diff();
  if (deconcat_dim_ == 0) {
    int offset_num = 0;
    for (int i = 0; i < top.size(); ++i) {
    	Blob<Dtype>* blob = top[i];
    	const Dtype* top_diff = blob->cpu_diff();
    	caffe_copy(blob->count(), top_diff,
    	           bottom_diff + (*bottom)[0]->offset(offset_num));
    	offset_num += blob->num();
    }
  }
}

#ifdef CPU_ONLY
STUB_GPU(DeConcatLayer);
#endif

INSTANTIATE_CLASS(DeConcatLayer);

}  // namespace caffe
