#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 GeneralizedSimilarityLossLayer<Dtype>::Forward_gpu(
    const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  int count1 = bottom[0]->count();
	int count2 = bottom[2]->count();
	const Dtype* bottom_data0=bottom[0]->gpu_data();
	const Dtype* bottom_data1=bottom[1]->gpu_data();
	const Dtype* bottom_data2=bottom[2]->gpu_data();
	const Dtype* bottom_data3=bottom[3]->gpu_data();
	const Dtype* bottom_data4=bottom[4]->gpu_data();
	const Dtype* bottom_data5=bottom[5]->gpu_data();
	caffe_gpu_mul(
	    count1,
		bottom_data0,
		bottom_data0,
		data_squ.mutable_gpu_data());
	caffe_gpu_gemv(
	      CblasNoTrans,
	      bottom[0]->num(),
	      bottom[0]->channels(),
	      Dtype(1.0),
		  data_squ.gpu_data(),  // a_i*a_i
	      summer_vec_.gpu_data(),
	      Dtype(0.0),
		  XA_norm.mutable_gpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1
	caffe_gpu_mul(
		    count1,
			bottom_data1,
			bottom_data1,
			data_squ.mutable_gpu_data());
		caffe_gpu_gemv(
		      CblasNoTrans,
		      bottom[0]->num(),
		      bottom[0]->channels(),
		      Dtype(1.0),
			  data_squ.gpu_data(),  // a_i*a_i
		      summer_vec_.gpu_data(),
		      Dtype(0.0),
			  YB_norm.mutable_gpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1
	caffe_gpu_mul(
		    count2,
			bottom_data2,
			bottom_data3,
			data23_squ.mutable_gpu_data());
	caffe_gpu_gemv(
		      CblasNoTrans,
		      bottom[2]->num(),
		      bottom[2]->channels(),
		      Dtype(1.0),
			  data23_squ.gpu_data(),  // a_i*a_i
		      summer_vec_2.gpu_data(),
		      Dtype(0.0),
			  data_XCYC.mutable_gpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1

Dtype alpha = 1/(2*bottom[0]->num());

	caffe_gpu_axpby(bottom[0]->num(), Dtype(1.0), XA_norm.gpu_data(), Dtype(0.0), similarity.mutable_gpu_data());
	caffe_gpu_axpby(bottom[0]->num(), Dtype(1.0), YB_norm.gpu_data(), Dtype(1.0), similarity.mutable_gpu_data());
	caffe_gpu_axpby(bottom[0]->num(), Dtype(-2.0), data_XCYC.gpu_data(), Dtype(1.0), similarity.mutable_gpu_data());
	caffe_gpu_axpby(bottom[0]->num(), Dtype(2.0), bottom_data4, Dtype(1.0), similarity.mutable_gpu_data());
	caffe_gpu_axpby(bottom[0]->num(), Dtype(2.0), bottom_data5, Dtype(1.0), similarity.mutable_gpu_data());

  Dtype margin = this->layer_param_.generalized_similarity_loss_param().margin();
  Dtype loss(0.0);
  Dtype averagedist=0;
  Dtype distthres=this->blobs_[0]->cpu_data()[0]*alpha;
  for (int i = 0; i < bottom[0]->num(); ++i) {
    averagedist+=similarity.cpu_data()[i];
    if (static_cast<int>(bottom[6]->cpu_data()[i])) {  // similar pairs
      loss += std::max(similarity.cpu_data()[i]-distthres+margin, Dtype(0.0));
    } else {  // dissimilar pairs
      loss += std::max(margin+distthres-similarity.cpu_data()[i], Dtype(0.0));
    }
  }
  //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;
  
  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;
    std::ofstream file0;
    std::ofstream file1;
    std::ofstream file2;
    std::ofstream file3;
    std::ofstream file4;
    std::ofstream file5;
    file.open(lossfile.data(), ios::out|ios::app);
    file0.open("./bottomcost3_0.txt", ios::out|ios::app);
    file1.open("./bottomcost3_1.txt", ios::out|ios::app);
    file2.open("./bottomcost3_2.txt", ios::out|ios::app);
    file3.open("./bottomcost3_3.txt", ios::out|ios::app);
    file4.open("./bottomcost3_4.txt", ios::out|ios::app);
    file5.open("./bottomcost3_5.txt", ios::out|ios::app);
      for (int i = 0; i < bottom[0]->num(); ++i)
      {
    	  file << i << "\t";
    	  file << similarity.cpu_data()[i] << "\n";
      }
      for (int i=0;i<bottom[0]->num(); ++i)
      {
          for (int j=0; j<bottom[0]->channels(); ++j) {
              const Dtype* bottom_i = bottom[0]->cpu_data()+bottom[0]->offset(i);
              file0 << bottom_i[j] << "\t";
          }
          file0 << "\n";
          for (int j=0; j<bottom[1]->channels(); ++j) {
              const Dtype* bottom_i = bottom[1]->cpu_data()+bottom[1]->offset(i);
              file1 << bottom_i[j] << "\t";
          }
          file1 << "\n";
          for (int j=0; j<bottom[2]->channels(); ++j) {
              const Dtype* bottom_i = bottom[2]->cpu_data()+bottom[2]->offset(i);
              file2 << bottom_i[j] << "\t";
          }
          file2 << "\n";
          for (int j=0; j<bottom[3]->channels(); ++j) {
              const Dtype* bottom_i = bottom[3]->cpu_data()+bottom[3]->offset(i);
              file3 << bottom_i[j] << "\t";
          }
          file3 << "\n";
          for (int j=0; j<bottom[4]->channels(); ++j) {
              const Dtype* bottom_i = bottom[4]->cpu_data()+bottom[4]->offset(i);
              file4 << bottom_i[j] << "\t";
          }
          file4 << "\n";
          for (int j=0; j<bottom[5]->channels(); ++j) {
              const Dtype* bottom_i = bottom[5]->cpu_data()+bottom[5]->offset(i);
              file5 << bottom_i[j] << "\t";
          }
          file5 << "\n";
      }
      file.close();
      file0.close();
      file1.close();
      file2.close();
      file3.close();
      file4.close();
      file5.close();
  }
  }
}

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

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

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


template <typename Dtype>
__global__ void GSLLTForward(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 GeneralizedSimilarityLossLayer<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_.generalized_similarity_loss_param().margin();
    caffe_gpu_set(num, Dtype(1.0), sum_num.mutable_gpu_data());
    Dtype* bottom_diff0=(*bottom)[0]->mutable_gpu_diff();
  Dtype* bottom_diff1=(*bottom)[1]->mutable_gpu_diff();
  Dtype* bottom_diff2=(*bottom)[2]->mutable_gpu_diff();
  Dtype* bottom_diff3=(*bottom)[3]->mutable_gpu_diff();
  Dtype* bottom_diff4=(*bottom)[4]->mutable_gpu_diff();
  Dtype* bottom_diff5=(*bottom)[5]->mutable_gpu_diff();
  const Dtype alpha = top[0]->cpu_diff()[0] / static_cast<Dtype>((*bottom)[0]->num());
    //XA and YB
          const int count01 = (*bottom)[0]->count();
      const int channels01 = (*bottom)[0]->channels();
    // NOLINT_NEXT_LINE(whitespace/operators)
      GSLLABForward<Dtype><<<CAFFE_GET_BLOCKS(count01), CAFFE_CUDA_NUM_THREADS>>>(
          count01, channels01, margin, distthres, alpha, 
          (*bottom)[6]->gpu_data(),  // pair similarity 0 or 1
          (*bottom)[0]->gpu_data(), (*bottom)[1]->gpu_data(),
          similarity.gpu_data(),  // the cached square distance between a and b
          bottom_diff0, bottom_diff1);
      CUDA_POST_KERNEL_CHECK;
      
      //XC and YC
      const int count23 = (*bottom)[2]->count();
      const int channels23 = (*bottom)[2]->channels();
      // NOLINT_NEXT_LINE(whitespace/operators)
      GSLLCForward<Dtype><<<CAFFE_GET_BLOCKS(count23), CAFFE_CUDA_NUM_THREADS>>>(
          count23, channels23, margin, distthres, alpha, 
          (*bottom)[6]->gpu_data(),  // pair similarity 0 or 1
          (*bottom)[2]->gpu_data(), (*bottom)[3]->gpu_data(),
          similarity.gpu_data(),  // the cached square distance between a and b
          bottom_diff2, bottom_diff3);
      CUDA_POST_KERNEL_CHECK;
      
      //XD and YE
      const int count45 = (*bottom)[4]->count();
      const int channels45 = (*bottom)[4]->channels();
      // NOLINT_NEXT_LINE(whitespace/operators)
      GSLLDEForward<Dtype><<<CAFFE_GET_BLOCKS(count45), CAFFE_CUDA_NUM_THREADS>>>(
          count45, channels45, margin, distthres, alpha, 
          (*bottom)[6]->gpu_data(),  // pair similarity 0 or 1
          similarity.gpu_data(),  // the cached square distance between a and b
          bottom_diff4, bottom_diff5);
      CUDA_POST_KERNEL_CHECK;
  //LOG(INFO) << "start blob";
  // NOLINT_NEXT_LINE(whitespace/operators)
      GSLLTForward<Dtype><<<CAFFE_GET_BLOCKS(num), CAFFE_CUDA_NUM_THREADS>>>(
          num, margin, distthres, top[0]->cpu_diff()[0], 
          (*bottom)[6]->gpu_data(),  // pair similarity 0 or 1
          similarity.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";
      if (this->layer_param_.has_recording()) {
  if (this->layer_param_.recording()) {
    std::ofstream file0;
    std::ofstream file1;
    std::ofstream file2;
    std::ofstream file3;
    std::ofstream file4;
    std::ofstream file5;
    file0.open("./bottomcost3_diff0.txt", ios::out|ios::app);
    file1.open("./bottomcost3_diff1.txt", ios::out|ios::app);
    file2.open("./bottomcost3_diff2.txt", ios::out|ios::app);
    file3.open("./bottomcost3_diff3.txt", ios::out|ios::app);
    file4.open("./bottomcost3_diff4.txt", ios::out|ios::app);
    file5.open("./bottomcost3_diff5.txt", ios::out|ios::app);
      for (int i=0;i<(*bottom)[0]->num(); ++i)
      {
          for (int j=0; j<(*bottom)[0]->channels(); ++j) {
              const Dtype* bottom_i = (*bottom)[0]->cpu_diff()+(*bottom)[0]->offset(i);
              file0 << bottom_i[j] << "\t";
          }
          file0 << "\n";
          for (int j=0; j<(*bottom)[1]->channels(); ++j) {
              const Dtype* bottom_i = (*bottom)[1]->cpu_diff()+(*bottom)[1]->offset(i);
              file1 << bottom_i[j] << "\t";
          }
          file1 << "\n";
          for (int j=0; j<(*bottom)[2]->channels(); ++j) {
              const Dtype* bottom_i = (*bottom)[2]->cpu_diff()+(*bottom)[2]->offset(i);
              file2 << bottom_i[j] << "\t";
          }
          file2 << "\n";
          for (int j=0; j<(*bottom)[3]->channels(); ++j) {
              const Dtype* bottom_i = (*bottom)[3]->cpu_diff()+(*bottom)[3]->offset(i);
              file3 << bottom_i[j] << "\t";
          }
          file3 << "\n";
          for (int j=0; j<(*bottom)[4]->channels(); ++j) {
              const Dtype* bottom_i = (*bottom)[4]->cpu_diff()+(*bottom)[4]->offset(i);
              file4 << bottom_i[j] << "\t";
          }
          file4 << "\n";
          for (int j=0; j<(*bottom)[5]->channels(); ++j) {
              const Dtype* bottom_i = (*bottom)[5]->cpu_diff()+(*bottom)[5]->offset(i);
              file5 << bottom_i[j] << "\t";
          }
          file5 << "\n";
      }
      file0.close();
      file1.close();
      file2.close();
      file3.close();
      file4.close();
      file5.close();
  }
  }
}

INSTANTIATE_CLASS(GeneralizedSimilarityLossLayer);

}  // namespace caffe
