#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 ConsNNPairLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
  int num=bottom[0]->num();
	Dtype* pairlabel=(*top)[2]->mutable_cpu_data();
	Dtype* data1idxarray=(*top)[0]->mutable_cpu_data();
	Dtype* data2idxarray=(*top)[1]->mutable_cpu_data();
	int dim=bottom[0]->count()/bottom[0]->num();
	const int npos=NNpos;
		const int nneg=NNneg;
		int nearneighbor_pos[npos];
		int nearneighbor_neg[nneg];
		Dtype nndist_pos[npos];
		Dtype nndist_neg[nneg];
	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;
	    int pairidx1=0;
	    int pairidx2=0;
	    for (int i=0; i<num; i++) {
	      for (int j=0;j<NNpos;j++) {nearneighbor_pos[j]=-1; nndist_pos[j]=DBL_MAX;}
	      for (int j=0;j<NNneg;j++) {nearneighbor_neg[j]=-1; nndist_neg[j]=DBL_MAX;}
	      Dtype* dist_i=dist_data+dist_mat_.offset(i);
	      for (int k=0; k<num; k++) {
	      if (i!=k) {
	        if (bottom_label[i]==bottom_label[k]) {
	        	if (dist_i[k]<nndist_pos[NNpos-1]) {
	        		          for (int j=0; j<NNpos; j++) {
	        		            if (dist_i[k]<nndist_pos[j]) {
	        		              for (int l=NNpos-1; l>j; l--) {
	        		                nndist_pos[l]=nndist_pos[l-1];
	        		                nearneighbor_pos[l]=nearneighbor_pos[l-1];
	        		              }
	        		              nndist_pos[j]=dist_i[k];
	        		              nearneighbor_pos[j]=k;
	        		              break;
	        		            }
	        		          }
	        		        }
	        } else {
	        if (dist_i[k]<nndist_neg[NNneg-1]) {
	          for (int j=0; j<NNneg; j++) {
	            if (dist_i[k]<nndist_neg[j]) {
	              for (int l=NNneg-1; l>j; l--) {
	                nndist_neg[l]=nndist_neg[l-1];
	                nearneighbor_neg[l]=nearneighbor_neg[l-1];
	              }
	              nndist_neg[j]=dist_i[k];
	              nearneighbor_neg[j]=k;
	              break;
	            }
	          }
	        }
	        }
	      }
	      }
	      for (int j=0;j<NNpos;j++) {
	    	  if (nearneighbor_pos[j]>=0) {
	    		  data1idxarray[pairidx]=i;
	    		  data2idxarray[pairidx]=nearneighbor_pos[j];
	    		  pairlabel[pairidx]=Dtype(1.0);
	    		  pairidx++;
	    		  pairidx1++;
	    	  }
	      }
	      for (int j=0;j<NNneg;j++) {
	      	    	  if (nearneighbor_neg[j]>=0) {
	      	    		  data1idxarray[pairidx]=i;
	      	    		  data2idxarray[pairidx]=nearneighbor_neg[j];
	      	    		  pairlabel[pairidx]=Dtype(0.0);
	      	    		  pairidx++;
	      	    		  pairidx2++;
	      	    	  }
	      }
	    }
	    for (int j=pairidx;j<pospairnum+negpairnum;j++)
	    {
	    	data1idxarray[j]=data1idxarray[pairidx];
	    	data2idxarray[j]=data2idxarray[pairidx];
	    	pairlabel[j]=pairlabel[pairidx];
	    }
	    /*LOG(INFO) << pairidx;
	    for (int i=0;i<6;i++) {
	     LOG(INFO) << data1idxarray[i] << "," << data2idxarray[i] << "," << pairlabel[i];
	    }*/
    //CHECK_EQ(pairidx,pospairnum+negpairnum);

}


INSTANTIATE_CLASS(ConsNNPairLayer);


}  // namespace caffe