#include <vector>

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

namespace caffe {

template <typename Dtype>
void MinusLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  CHECK_EQ(bottom[0]->num(), bottom[1]->num()) << "num should be equal";
  CHECK_EQ(bottom[0]->channels(), bottom[1]->channels()) << "channels should be equal";
  CHECK_EQ(bottom[0]->height(), bottom[1]->height()) << "height should be equal";
  CHECK_EQ(bottom[0]->width(), bottom[1]->width()) << "width should be equal";
}

template <typename Dtype>
void MinusLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  // Initialize with the first blob.
  int count_ = bottom[0]->count();
  int num_ = bottom[0]->num();
  int channels_ = bottom[0]->channels();
  int height_ = bottom[0]->height();
  int width_ = bottom[0]->width();

  (*top)[0]->Reshape(num_, channels_, height_, width_);

}

template <typename Dtype>
void MinusLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  Dtype* top_data = (*top)[0]->mutable_cpu_data();
  int count = bottom[0]->count();
  caffe_sub(
      count,
      bottom[0]->cpu_data(),  // a
      bottom[1]->cpu_data(),  // b
      top_data);  // a_i-b_i
}

template <typename Dtype>
void MinusLayer<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* bottom1_diff = (*bottom)[0]->mutable_cpu_diff();
  Dtype* bottom2_diff = (*bottom)[1]->mutable_cpu_diff();
  int count = (*bottom)[0]->count();
  caffe_copy(count, top_diff, bottom1_diff);
  caffe_copy(count, top_diff, bottom2_diff);
  caffe_scal(count, Dtype(-1.0), bottom2_diff);
}

#ifdef CPU_ONLY
STUB_GPU(MinusLayer);
#endif

INSTANTIATE_CLASS(MinusLayer);

}  // namespace caffe
