#include <algorithm>
#include <cmath>
#include <cfloat>
#include <vector>

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

namespace caffe {


template <typename Dtype>
void ConsPairLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  int num=bottom[0]->num();
	Dtype* pairlabel=(*top)[2]->mutable_cpu_data();
	const int NN=10;
	Dtype* data1idxarray=(*top)[0]->mutable_cpu_data();
	Dtype* data2idxarray=(*top)[1]->mutable_cpu_data();
	int dim=bottom[0]->count()/bottom[0]->num();
	int nearneighbor[NN];
	Dtype nndist[NN];
	const Dtype* bottom_data=bottom[0]->gpu_data();
	const Dtype* bottom_label=bottom[1]->cpu_data();//default: 1~maxlabel
	const int count = bottom[0]->count();
	caffe_gpu_mul(
      count,
      bottom[0]->gpu_data(),  // a
      bottom[0]->gpu_data(),  // a
      temp_mat_.mutable_gpu_data());  // a_i*a_i
    caffe_gpu_gemv(
      CblasNoTrans,
      bottom[0]->num(), bottom[0]->channels(),
      Dtype(1.0),
      temp_mat_.gpu_data(),  // a_i*a_i
      summer_vec_.gpu_data(),
      Dtype(0.0),
      temp_vec_.mutable_gpu_data());  // \Sum a_i*a_i
    caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[0]->num(), bottom[0]->num(), 1,
      Dtype(1.0),
      temp_vec_.gpu_data(), summer_vec_num_.gpu_data(),
      Dtype(0),
      dist_mat_.mutable_gpu_data());
    caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[0]->num(), bottom[0]->num(), 1,
      Dtype(1.0),
      summer_vec_num_.gpu_data(), temp_vec_.gpu_data(),
      Dtype(1.0),
      dist_mat_.mutable_gpu_data());
    caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[0]->num(), bottom[0]->num(), bottom[0]->channels(),
      Dtype(-2.0),
      bottom[0]->gpu_data(), bottom[0]->gpu_data(),
      Dtype(1.0),
      dist_mat_.mutable_gpu_data());
      Dtype* dist_data=dist_mat_.mutable_cpu_data();
    int pairidx=0;
    for (int i=0; i<num; i++) {
      for (int j=0;j<NN;j++) {nearneighbor[j]=-1; nndist[j]=DBL_MAX;}
      Dtype* dist_i=dist_data+dist_mat_.offset(i);
      for (int k=0; k<num; k++) {
        if (bottom_label[i]==bottom_label[k]) {
          if (i<k) {data1idxarray[pairidx]=i;data2idxarray[pairidx]=k; pairlabel[pairidx]=Dtype(1.0); pairidx++;}
          continue;
        }
        if (dist_i[k]<nndist[NN-1]) {
          for (int j=0; j<NN; j++) {
            if (dist_i[k]<nndist[j]) {
              for (int l=NN-1; l>j; l--) {
                nndist[l]=nndist[l-1];
                nearneighbor[l]=nearneighbor[l-1];
              }
              nndist[j]=dist_i[k];
              nearneighbor[j]=k;
		break;
            }
          }
        }
      }
      for (int j=0;j<NN;j++) {
        if (nearneighbor[j]>=0) {data1idxarray[pairidx]=i;data2idxarray[pairidx]=nearneighbor[j]; pairlabel[pairidx]=Dtype(0.0); pairidx++;}
      }
    }
    CHECK_EQ(pairidx,pospairnum+negpairnum);

}


INSTANTIATE_CLASS(ConsPairLayer);


}  // namespace caffe