#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 DoubletLossLayer<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_.mutable_gpu_data());  // a_i-b_i
  caffe_gpu_powx(
      count,
      diff_.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_sq_.mutable_gpu_data());  // \Sum (a_i-b_i)^2
  Dtype margin = this->layer_param_.doublet_loss_param().margin();
  Dtype distthres=this->blobs_[0]->cpu_data()[0];
  Dtype loss(0.0);
  for (int i = 0; i < bottom[0]->num(); ++i) {
    if (static_cast<int>(bottom[2]->cpu_data()[i])) {  // similar pairs
      loss += std::max(dist_sq_.cpu_data()[i]-distthres+margin, Dtype(0.0));
    } else {  // dissimilar pairs
      loss += std::max(margin+distthres-dist_sq_.cpu_data()[i], Dtype(0.0));
    }
  }
  loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  (*top)[0]->mutable_cpu_data()[0] = loss;
  //LOG(INFO) << this->layer_param_.has_recording() << "," << this->layer_param_.recording();
  /*Dtype meandist=0;
  Dtype x1=0;
  Dtype x2=0;
  for (int i = 0; i < bottom[0]->num(); ++i)
  {
  meandist=meandist+dist_sq_.cpu_data()[i];
  }
  LOG(INFO) << distthres << "," << meandist/bottom[0]->num();
  for (int i = 0; i < bottom[0]->count(); ++i)
  {
  x1=x1+fabs(bottom[0]->cpu_data()[i]);
  x2=x2+fabs(bottom[1]->cpu_data()[i]);
  }
  LOG(INFO) << x1/bottom[0]->num() << "," << x2/bottom[0]->num();*/
  /*for (int i = 0; i < 20; ++i)
  {
    LOG(INFO) << dist_sq_.cpu_data()[i] << "\t" << bottom[2]->cpu_data()[i];
  }*/
  string lossfile(this->layer_param_.name());
  lossfile="./"+lossfile+".txt";
/*  if (this->layer_param_.has_recording()) {
  if (this->layer_param_.recording()) {
  //LOG(INFO) << "cl 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();
  }
  }
  */
  Dtype* dist_data=dist_sq_.mutable_cpu_data();
    if (this->layer_param_.has_recording()) {
  if (this->layer_param_.recording()) {
  //LOG(INFO) << "cl recording";
    std::ofstream file;
    file.open(lossfile.data(), ios::out|ios::app);
      for (int i = 0; i < bottom[0]->num(); ++i)
      {
    	  Dtype* dist_i=dist_data+dist_sq_.offset(i);
    	  for (int j=0;j<bottom[1]->num();j++) {
    	      file << dist_i[j] << "\t";
    	  }
    	  file << "\n";
      }
      file.close();
  }
  }
}

template <typename Dtype>
__global__ void DLLForward(const int count, const int channels,
    const Dtype margin, const Dtype distthres, const Dtype alpha, 
    const Dtype* y, const Dtype* diff, const Dtype* dist_sq,
    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 ((dist_sq[n]-distthres+margin) > 0.0) {
        bottom_diff[i] = alpha * diff[i];
        //dt_diff[0] = dt_diff[0] - 0.25 * topdiff/bottomnum;
      } else {
        bottom_diff[i] = 0;
      }
    } else {  // dissimilar pairs
      if ((margin+distthres-dist_sq[n]) > 0.0) {
        bottom_diff[i] = -alpha * diff[i];
        //dt_diff[0] = dt_diff[0] + 0.25 * topdiff/bottomnum;
      } else {
        bottom_diff[i] = 0;
      }
    }
  }
}

template <typename Dtype>
__global__ void DLLTForward(const int num, const Dtype margin, const Dtype distthres, const Dtype topdiff,
    const Dtype* y, const Dtype* dist_sq, Dtype *dt_diff) {
  CUDA_KERNEL_LOOP(n, num) {
    if (static_cast<int>(y[n])) {  // similar pairs
      if ((dist_sq[n]-distthres+margin) > 0.0) {
        dt_diff[n] =  - 0.5 * topdiff/num;
      } else {
        dt_diff[n] = 0;
      }
    } else {  // dissimilar pairs
      if ((margin+distthres-dist_sq[n]) > 0.0) {
        dt_diff[n] = 0.5 * topdiff/num;
      } else {
        dt_diff[n] = 0;
      }
    }
  }
}









template <typename Dtype>
void DoubletLossLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
    Dtype distthres=this->blobs_[0]->cpu_data()[0];
    this->blobs_[0]->mutable_cpu_diff()[0]=0;
    Blob<Dtype> dtdiff, sum_num;
    int num=(*bottom)[0]->num();
    dtdiff.Reshape(num, 1, 1, 1);
    sum_num.Reshape(num, 1, 1, 1);
    Dtype margin = this->layer_param_.doublet_loss_param().margin();
    caffe_gpu_set(num, Dtype(1.0), sum_num.mutable_gpu_data());
  for (int i = 0; i < 2; ++i) {
    if (propagate_down[i]) {
      const int count = (*bottom)[0]->count();
      const int channels = (*bottom)[0]->channels();
      const Dtype sign = (i == 0) ? 1 : -1;
      const Dtype alpha = sign * top[0]->cpu_diff()[0] /
          static_cast<Dtype>((*bottom)[0]->num());
          //LOG(INFO) << "start" << i;
      // NOLINT_NEXT_LINE(whitespace/operators)
      DLLForward<Dtype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
          count, channels, margin, distthres, alpha, 
          (*bottom)[2]->gpu_data(),  // pair similarity 0 or 1
          diff_.gpu_data(),  // the cached eltwise difference between a and b
          dist_sq_.gpu_data(),  // the cached square distance between a and b
          (*bottom)[i]->mutable_gpu_diff());
      CUDA_POST_KERNEL_CHECK;
      //LOG(INFO) << "end" << i;
    }
  }
  //LOG(INFO) << "start blob";
  // NOLINT_NEXT_LINE(whitespace/operators)
      DLLTForward<Dtype><<<CAFFE_GET_BLOCKS(num), CAFFE_CUDA_NUM_THREADS>>>(
          num, margin, distthres, top[0]->cpu_diff()[0], 
          (*bottom)[2]->gpu_data(),  // pair similarity 0 or 1
          dist_sq_.gpu_data(),  // the cached square distance between a and b
          dtdiff.mutable_gpu_data());
      CUDA_POST_KERNEL_CHECK;
      //LOG(INFO) << "end blob";
      Dtype dot(0);
      caffe_gpu_dot(num, dtdiff.gpu_data(), sum_num.gpu_data(), &dot);
      this->blobs_[0]->mutable_cpu_diff()[0]=dot;
      //LOG(INFO) << "end dot";
}

INSTANTIATE_CLASS(DoubletLossLayer);

}  // namespace caffe
