#include <algorithm>
#include <vector>

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

namespace caffe {

template <typename Dtype>
void TripletLossLayer<Dtype>::Forward_gpu(
    const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  const int count = bottom[0]->count();
  caffe_gpu_sub(
      count,
      bottom[0]->gpu_data(),  // a
      bottom[1]->gpu_data(),  // b
      diff_ij_.mutable_gpu_data());  // a_i-b_i
  caffe_gpu_sub(
      count,
      bottom[0]->gpu_data(),  // a
      bottom[2]->gpu_data(),  // b
      diff_ik_.mutable_gpu_data());  // a_i-b_i
  caffe_gpu_powx(
      count,
      diff_ij_.mutable_gpu_data(),  // a_i-b_i
      Dtype(2),
      diff_sq_.mutable_gpu_data());  // (a_i-b_i)^2
  caffe_gpu_gemv(
      CblasNoTrans,
      bottom[0]->num(),
      bottom[0]->channels(),
      Dtype(1.0),
      diff_sq_.gpu_data(),  // (a_i-b_i)^2
      summer_vec_.gpu_data(),
      Dtype(0.0),
      dist_ij_.mutable_gpu_data());  // \Sum (a_i-b_i)^2
  caffe_gpu_powx(
      count,
      diff_ik_.mutable_gpu_data(),  // a_i-b_i
      Dtype(2),
      diff_sq_.mutable_gpu_data());  // (a_i-b_i)^2
  caffe_gpu_gemv(
      CblasNoTrans,
      bottom[0]->num(),
      bottom[0]->channels(),
      Dtype(1.0),
      diff_sq_.gpu_data(),  // (a_i-b_i)^2
      summer_vec_.gpu_data(),
      Dtype(0.0),
      dist_ik_.mutable_gpu_data());  // \Sum (a_i-b_i)^2
      Dtype margin = this->layer_param_.contrastive_loss_param().margin();
  Dtype loss(0.0);
  for (int i = 0; i < bottom[0]->num(); ++i) {
	  loss += std::max(margin - dist_ik_.cpu_data()[i] + dist_ij_.cpu_data()[i], Dtype(0.0));
	  //loss += dist_ij_.cpu_data()[i];
  }
  loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  (*top)[0]->mutable_cpu_data()[0] = loss;
}

template <typename Dtype>
__global__ void TLLForward(const int count, const int channels, const Dtype margin, const Dtype alpha,
    const Dtype* diff_ij_, const Dtype* diff_ik_, const Dtype* dist_ij_, const Dtype* dist_ik_,
    Dtype *bottom_diff0, Dtype *bottom_diff1, Dtype *bottom_diff2) {
  CUDA_KERNEL_LOOP(i, count) {
    int n = i / channels;  // the num index, to access y and dist_sq
    if (margin - dist_ik_[n]+dist_ij_[n]>0) {
      //bottom_diff0[i]=2*alpha*diff_ij_[i] - alpha*diff_ik_[i];
      //bottom_diff1[i]= (-2)*alpha*diff_ij_[i];
      //bottom_diff2[i]=2*alpha*diff_ik_[i];
      bottom_diff0[i]=alpha*diff_ij_[i] - alpha*diff_ik_[i];
      bottom_diff1[i]= (-1)*alpha*diff_ij_[i];
      bottom_diff2[i]=alpha*diff_ik_[i];
    } else {
      //bottom_diff0[i]=alpha*diff_ij_[i];
      //bottom_diff1[i]= (-1)*alpha*diff_ij_[i];
      //bottom_diff2[i]=0;
      bottom_diff0[i]=0;
      bottom_diff1[i]=0;
      bottom_diff2[i]=0;
    }
  }
}

template <typename Dtype>
void TripletLossLayer<Dtype>::Backward_gpu(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();
	const int count = (*bottom)[0]->count();
	Dtype margin = this->layer_param_.contrastive_loss_param().margin();
	Dtype* bottom_diff0=(*bottom)[0]->mutable_gpu_diff();
	Dtype* bottom_diff1=(*bottom)[1]->mutable_gpu_diff();
	Dtype* bottom_diff2=(*bottom)[2]->mutable_gpu_diff();
	TLLForward<Dtype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
          count, channels, margin, alpha,
          diff_ij_.gpu_data(),  // the cached eltwise difference between x_i and x_j
          diff_ik_.gpu_data(),  // the cached eltwise difference between x_i and x_k
          dist_ij_.gpu_data(),  // the cached square distance between x_i and x_j
          dist_ik_.gpu_data(),  // the cached square distance between x_i and x_k
          (*bottom)[0]->mutable_gpu_diff(), (*bottom)[1]->mutable_gpu_diff(), (*bottom)[2]->mutable_gpu_diff());
      CUDA_POST_KERNEL_CHECK;
}

INSTANTIATE_CLASS(TripletLossLayer);

}  // namespace caffe
