#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 TripletLossLayer<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[1]->channels(), bottom[2]->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]->height(), 1);
  CHECK_EQ(bottom[2]->width(), 1);
}

template <typename Dtype>
void TripletLossLayer<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_ij_.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
	  diff_ik_.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
	  diff_sq_.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
	  dist_ij_.Reshape(bottom[0]->num(), 1, 1, 1);
	  dist_ik_.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 TripletLossLayer<Dtype>::Forward_cpu(
    const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
  int count = bottom[0]->count();
  caffe_sub(
      count,
      bottom[0]->cpu_data(),  // a
      bottom[1]->cpu_data(),  // b
      diff_ij_.mutable_cpu_data());  // a_i-b_i
  const int channels = bottom[0]->channels();
  Dtype loss(0.0);
  for (int i = 0; i < bottom[0]->num(); ++i) {
    dist_ij_.mutable_cpu_data()[i] = caffe_cpu_dot(channels,
        diff_ij_.cpu_data() + (i*channels), diff_ij_.cpu_data() + (i*channels));
  }
  caffe_sub(
        count,
        bottom[0]->cpu_data(),  // a
        bottom[2]->cpu_data(),  // b
        diff_ik_.mutable_cpu_data());  // a_i-b_i
  for (int i = 0; i < bottom[0]->num(); ++i) {
	  dist_ik_.mutable_cpu_data()[i] = caffe_cpu_dot(channels,
	          diff_ik_.cpu_data() + (i*channels), diff_ik_.cpu_data() + (i*channels));
  }
  for (int i = 0; i < bottom[0]->num(); ++i) {
	  loss += std::max(Dtype(1.0) - dist_ik_.cpu_data()[i] + dist_ij_.cpu_data()[i], Dtype(0.0));
	  loss += dist_ij_.cpu_data()[i];
  }
  //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;
}

template <typename Dtype>
void TripletLossLayer<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();
	Dtype* bottom_diff0=(*bottom)[0]->mutable_cpu_diff();
	Dtype* bottom_diff1=(*bottom)[1]->mutable_cpu_diff();
	Dtype* bottom_diff2=(*bottom)[2]->mutable_cpu_diff();
	for (int j = 0; j < num; ++j) {
		if (Dtype(1.0) - dist_ik_.cpu_data()[j]+dist_ij_.cpu_data()[j]>0) {
			caffe_cpu_axpby(
					channels,
				              alpha*2,
				              diff_ij_.cpu_data() + (j*channels),
				              Dtype(0.0),
							  bottom_diff0 + (j*channels));
			caffe_cpu_axpby(
					channels,
					              -alpha,
					              diff_ik_.cpu_data() + (j*channels),
					              Dtype(1.0),
								  bottom_diff0 + (j*channels));
			caffe_cpu_axpby(
								channels,
							              alpha*(-2),
							              diff_ij_.cpu_data() + (j*channels),
							              Dtype(0.0),
										  bottom_diff1 + (j*channels));
			caffe_cpu_axpby(
								channels,
								              alpha,
								              diff_ik_.cpu_data() + (j*channels),
								              Dtype(0.0),
											  bottom_diff2 + (j*channels));
		} else {
			caffe_cpu_axpby(
					channels,
				              alpha,
				              diff_ij_.cpu_data() + (j*channels),
				              Dtype(0.0),
							  bottom_diff0 + (j*channels));
			caffe_cpu_axpby(
								channels,
							              -alpha,
							              diff_ij_.cpu_data() + (j*channels),
							              Dtype(0.0),
										  bottom_diff1 + (j*channels));
			caffe_set(channels, Dtype(0), bottom_diff2 + (j*channels));
		}
	}



}

#ifdef CPU_ONLY
STUB_GPU(TripletLossLayer);
#endif

INSTANTIATE_CLASS(TripletLossLayer);

}  // namespace caffe
