#include <algorithm>
#include <vector>

#include "caffe/layer.hpp"
#include "caffe/loss_layers.hpp"
#include "caffe/util/io.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

template <typename Dtype>
void LinearContrastiveLossLayer<Dtype>::LayerSetUp(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  LossLayer<Dtype>::LayerSetUp(bottom, top);
  CHECK_EQ(bottom[0]->channels(), 1);
  CHECK_EQ(bottom[0]->height(), 1);
  CHECK_EQ(bottom[0]->width(), 1);
  CHECK_EQ(bottom[1]->channels(), 1);
  CHECK_EQ(bottom[1]->height(), 1);
  CHECK_EQ(bottom[1]->width(), 1);
}

template <typename Dtype>
void LinearContrastiveLossLayer<Dtype>::Reshape(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
	LossLayer<Dtype>::Reshape(bottom, top);
	//LOG(INFO) << "layersetup:" << bottom[0]->num() << "," << bottom[0]->channels() << "," << bottom[0]->height() << "," << bottom[0]->width();
	  // vector of ones used to sum along channels
	  summer_vec_.Reshape(bottom[0]->channels(), 1, 1, 1);
	  for (int i = 0; i < bottom[0]->channels(); ++i)
	    summer_vec_.mutable_cpu_data()[i] = Dtype(1);
}

template <typename Dtype>
void LinearContrastiveLossLayer<Dtype>::Forward_cpu(
    const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
  int count = bottom[0]->count();
  Dtype loss(0.0);
  Dtype averagedist=0;
  for (int i = 0; i < bottom[0]->num(); ++i) {
    if (static_cast<int>(bottom[1]->cpu_data()[i])) {  // similar pairs
      loss += std::max(bottom[0]->cpu_data()[i]+Dtype(1.0), Dtype(0.0));
    } else {  // dissimilar pairs
      loss += std::max(-bottom[0]->cpu_data()[i]+Dtype(1.0), Dtype(0.0));
    }
  }
  //LOG(INFO) << "average distance:" << averagedist/bottom[0]->num();
  loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  (*top)[0]->mutable_cpu_data()[0] = loss;
  //LOG(INFO) << "contrastive_loss:" << loss;

}

template <typename Dtype>
void LinearContrastiveLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
      const Dtype alpha = top[0]->cpu_diff()[0] /
          static_cast<Dtype>((*bottom)[0]->num());
      int num = (*bottom)[0]->num();
      int channels = (*bottom)[0]->channels();
      for (int j = 0; j < num; ++j) {
        Dtype* bout = (*bottom)[0]->mutable_cpu_diff();
        if (static_cast<int>((*bottom)[1]->cpu_data()[j])) {  // similar pairs
          if ((*bottom)[0]->cpu_data()[j]>-1) {
        	  bout[j]=alpha * (*bottom)[1]->cpu_data()[j];
          } else {
        	  bout[j]=Dtype(0);
          }
        } else {  // dissimilar pairs
        	if ((*bottom)[0]->cpu_data()[j]<1) {
        	  bout[j]=-alpha * (1-(*bottom)[1]->cpu_data()[j]);
        	} else {
        	  bout[j]=Dtype(0);
        	}
        }
      }
}

#ifdef CPU_ONLY
STUB_GPU(LinearContrastiveLossLayer);
#endif

INSTANTIATE_CLASS(LinearContrastiveLossLayer);

}  // namespace caffe
