#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 LinearContrastiveLossLayer<Dtype>::Forward_gpu(
    const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  const int count = bottom[0]->count();
  
  Dtype loss(0.0);
  
  Dtype correctnum(0.0);
  
  Dtype margin = this->layer_param_.contrastive_loss_param().margin();
  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]+margin, Dtype(0.0));
      if (bottom[0]->cpu_data()[i]<0) {correctnum=correctnum+1;}
    } else {  // dissimilar pairs
      loss += std::max(-bottom[0]->cpu_data()[i]+margin, Dtype(0.0));
      if (bottom[0]->cpu_data()[i]>0) {correctnum=correctnum+1;}
    }
  }
  loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  //LOG(INFO) << "verification accuracy (CI):" << correctnum/bottom[0]->num();
  (*top)[0]->mutable_cpu_data()[0] = 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 << bottom[0]->cpu_data()[i] << "\n";
      }
      file.close();
      }
   }
}

template <typename Dtype>
__global__ void LCLLForward(const int count, const int channels,
    const Dtype alpha, const Dtype margin,
    const Dtype* y, const Dtype* x,
    Dtype *bottom_diff) {
  CUDA_KERNEL_LOOP(i, count) {
    int n = i / channels;  // the num index, to access y and dist_sq
    if (static_cast<int>(y[n])) {  // similar pairs
      if (x[n] > -margin) {
        bottom_diff[i] = alpha * y[n];
      } else {
        bottom_diff[i] = 0;
      }
    } else {  // dissimilar pairs
      if (x[n] < margin) {
        bottom_diff[i] = -alpha * (1.0 - y[n]);
      } else {
        bottom_diff[i] = 0;
      }
    }
  }
}

template <typename Dtype>
void LinearContrastiveLossLayer<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();
      const Dtype alpha = top[0]->cpu_diff()[0] /
          static_cast<Dtype>((*bottom)[0]->num());
          Dtype margin = this->layer_param_.contrastive_loss_param().margin();
      // NOLINT_NEXT_LINE(whitespace/operators)
      LCLLForward<Dtype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
          count, channels, alpha, margin,
          (*bottom)[1]->gpu_data(),  // pair similarity 0 or 1
          (*bottom)[0]->gpu_data(),
          (*bottom)[0]->mutable_gpu_diff());
      CUDA_POST_KERNEL_CHECK;
}

INSTANTIATE_CLASS(LinearContrastiveLossLayer);

}  // namespace caffe
