#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 LinearTripletLossLayer<Dtype>::Forward_gpu(
    const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  const int count = bottom[0]->count();
  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+bottom[0]->cpu_data()[i] - bottom[1]->cpu_data()[i], Dtype(0.0));
  }
  loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  (*top)[0]->mutable_cpu_data()[0] = loss;
  //******************
  /*std::ofstream file;
    file.open("./dist.txt", 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>
__global__ void LTLLForward(const int count, const int channels, const Dtype margin,
    const Dtype alpha,
    const Dtype* t1, const Dtype* t2,
    Dtype *bottom_diff0, Dtype *bottom_diff1) {
  CUDA_KERNEL_LOOP(i, count) {
    int n = i / channels;  // the num index, to access y and dist_sq
    if (margin + t1[n] - t2[n]>0) {
        	bottom_diff0[n] = alpha;
        	bottom_diff1[n] = -alpha;
        } else {
        	bottom_diff0[n] = Dtype(0.0);
        	bottom_diff1[n] = Dtype(0.0);
        }
  }
}

template <typename Dtype>
void LinearTripletLossLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
      const int count = (*bottom)[0]->count();
      const int channels = (*bottom)[0]->channels();
      Dtype margin = this->layer_param_.contrastive_loss_param().margin();
      const Dtype alpha = top[0]->cpu_diff()[0] /
          static_cast<Dtype>((*bottom)[0]->num());
      // NOLINT_NEXT_LINE(whitespace/operators)
      LTLLForward<Dtype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
          count, channels, margin, alpha,
          (*bottom)[0]->gpu_data(), 
          (*bottom)[1]->gpu_data(),
          (*bottom)[0]->mutable_gpu_diff(),
          (*bottom)[1]->mutable_gpu_diff());
      CUDA_POST_KERNEL_CHECK;
}

INSTANTIATE_CLASS(LinearTripletLossLayer);

}  // namespace caffe
