#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 ContrastiveLossLayer<Dtype>::LayerSetUp(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  LossLayer<Dtype>::LayerSetUp(bottom, top);
  CHECK_EQ(bottom[0]->channels(), bottom[1]->channels());
  CHECK_EQ(bottom[0]->height(), 1);
  CHECK_EQ(bottom[0]->width(), 1);
  CHECK_EQ(bottom[1]->height(), 1);
  CHECK_EQ(bottom[1]->width(), 1);
  CHECK_EQ(bottom[2]->channels(), 1);
  CHECK_EQ(bottom[2]->height(), 1);
  CHECK_EQ(bottom[2]->width(), 1);
}

template <typename Dtype>
void ContrastiveLossLayer<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();
	  diff_.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
	  diff_sq_.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
	  dist_sq_.Reshape(bottom[0]->num(), 1, 1, 1);
	  // 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 ContrastiveLossLayer<Dtype>::Forward_cpu(
    const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
  int count = bottom[0]->count();
  Dtype* diff_data = diff_.mutable_cpu_data();
  caffe_sub(
      count,
      bottom[0]->cpu_data(),  // a
      bottom[1]->cpu_data(),  // b
      diff_.mutable_cpu_data());  // a_i-b_i
  const int channels = bottom[0]->channels();
  Dtype margin = this->layer_param_.contrastive_loss_param().margin();
  Dtype loss(0.0);
  Dtype averagedist=0;
  for (int i = 0; i < bottom[0]->num(); ++i) {
    dist_sq_.mutable_cpu_data()[i] = caffe_cpu_dot(channels,
        diff_.cpu_data() + (i*channels), diff_.cpu_data() + (i*channels));
    averagedist+=dist_sq_.cpu_data()[i];
    if (static_cast<int>(bottom[2]->cpu_data()[i])) {  // similar pairs
      loss += dist_sq_.cpu_data()[i];
    } else {  // dissimilar pairs
      loss += std::max(margin-dist_sq_.cpu_data()[i], 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;
  string lossfile(this->layer_param_.name());
  lossfile="./"+lossfile+".txt";
  if (this->layer_param_.has_recording()) {
    if (this->layer_param_.recording()) {
      std::ofstream file;
      file.open(lossfile.data(), ios::out|ios::app);
        for (int i = 0; i < bottom[0]->num(); ++i)
        {
      	  file << i << "\t";
      	  file << dist_sq_.cpu_data()[i] << "\n";
        }
        file.close();
    }
  }
}

template <typename Dtype>
void ContrastiveLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
  Dtype margin = this->layer_param_.contrastive_loss_param().margin();
  for (int i = 0; i < 2; ++i) {
    if (propagate_down[i]) {
      const Dtype sign = (i == 0) ? 1 : -1;
      const Dtype alpha = sign * top[0]->cpu_diff()[0] /
          static_cast<Dtype>((*bottom)[i]->num());
      int num = (*bottom)[i]->num();
      int channels = (*bottom)[i]->channels();
      for (int j = 0; j < num; ++j) {
        Dtype* bout = (*bottom)[i]->mutable_cpu_diff();
        if (static_cast<int>((*bottom)[2]->cpu_data()[j])) {  // similar pairs
          caffe_cpu_axpby(
              channels,
              alpha,
              diff_.cpu_data() + (j*channels),
              Dtype(0.0),
              bout + (j*channels));
        } else {  // dissimilar pairs
          if ((margin-dist_sq_.cpu_data()[j]) > Dtype(0.0)) {
            caffe_cpu_axpby(
                channels,
                -alpha,
                diff_.cpu_data() + (j*channels),
                Dtype(0.0),
                bout + (j*channels));
          } else {
            caffe_set(channels, Dtype(0), bout + (j*channels));
          }
        }
      }
    }
  }
}

#ifdef CPU_ONLY
STUB_GPU(ContrastiveLossLayer);
#endif

INSTANTIATE_CLASS(ContrastiveLossLayer);

}  // namespace caffe
