// Copyright 2014 BVLC and contributors.

#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"

using std::max;

namespace caffe {

template <typename Dtype>
void ConsNNPairLayer<Dtype>::LayerSetUp(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
	CHECK_EQ(bottom[0]->num(), bottom[1]->num())
		      << "The data and label should have the same number.";
}

template <typename Dtype>
void ConsNNPairLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
    int maxlabel=-1;
    int minlabel=100000;
    const Dtype* data=bottom[0]->cpu_data();
    const Dtype* label=bottom[1]->cpu_data();//default: 1~maxlabel
    //ConsNNPairParameter cons_nn_pair_param = this->layer_param_.cons_nn_pair_param_();
    NNpos=this->layer_param_.cons_nn_pair_param().nn_positive();
    NNneg=this->layer_param_.cons_nn_pair_param().nn_negative();
    temp_mat_.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
    summer_vec_.Reshape(bottom[0]->channels(), 1, 1, 1);
    for (int i = 0; i < bottom[0]->channels(); ++i)
    	    summer_vec_.mutable_cpu_data()[i] = Dtype(1);
    summer_vec_num_.Reshape(bottom[0]->num(), 1, 1, 1);
    for (int i = 0; i < bottom[0]->num(); ++i)
        	summer_vec_num_.mutable_cpu_data()[i] = Dtype(1);
    temp_vec_.Reshape(bottom[0]->num(), 1, 1, 1);
    dist_mat_.Reshape(bottom[0]->num(), bottom[0]->num(), 1, 1);
    pospairnum=bottom[0]->num()*NNpos;
    negpairnum=bottom[0]->num()*NNneg;
    //LOG(INFO) << "pos:" << pospairnum << ", neg:" << negpairnum;
    (*top)[0]->Reshape(pospairnum+negpairnum, 1, 1, 1);//data1
    (*top)[1]->Reshape(pospairnum+negpairnum, 1, 1, 1);//data2
    (*top)[2]->Reshape(pospairnum+negpairnum, 1, 1, 1);//pair label:pos(1),neg(0)
}

template <typename Dtype>
void ConsNNPairLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
	  int num=bottom[0]->num();
		Dtype* pairlabel=(*top)[2]->mutable_cpu_data();
		int dim=bottom[0]->count()/bottom[0]->num();
		Dtype* data1idxarray=(*top)[0]->mutable_cpu_data();
		Dtype* data2idxarray=(*top)[1]->mutable_cpu_data();
		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]->cpu_data();
		const Dtype* bottom_label=bottom[1]->cpu_data();//default: 1~maxlabel
		const int count = bottom[0]->count();
		caffe_mul(
	      count,
	      bottom[0]->cpu_data(),  // a
	      bottom[0]->cpu_data(),  // a
	      temp_mat_.mutable_cpu_data());  // a_i*a_i
	    caffe_cpu_gemv(
	      CblasNoTrans,
	      bottom[0]->num(), bottom[0]->channels(),
	      Dtype(1.0),
	      temp_mat_.cpu_data(),  // a_i*a_i
	      summer_vec_.cpu_data(),
	      Dtype(0.0),
	      temp_vec_.mutable_cpu_data());  // \Sum a_i*a_i
	    caffe_cpu_gemm(
	      CblasNoTrans, CblasTrans,
	      bottom[0]->num(), bottom[0]->num(), 1,
	      Dtype(1.0),
	      temp_vec_.cpu_data(), summer_vec_num_.cpu_data(),
	      Dtype(0),
	      dist_mat_.mutable_cpu_data());
	    caffe_cpu_gemm(
	      CblasNoTrans, CblasTrans,
	      bottom[0]->num(), bottom[0]->num(), 1,
	      Dtype(1.0),
	      summer_vec_num_.cpu_data(), temp_vec_.cpu_data(),
	      Dtype(1.0),
	      dist_mat_.mutable_cpu_data());
	    caffe_cpu_gemm(
	      CblasNoTrans, CblasTrans,
	      bottom[0]->num(), bottom[0]->num(), bottom[0]->channels(),
	      Dtype(-2.0),
	      bottom[0]->cpu_data(), bottom[0]->cpu_data(),
	      Dtype(1.0),
	      dist_mat_.mutable_cpu_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) << pairidx1 << "," << pairidx2;
	    //CHECK_EQ(pairidx,pospairnum+negpairnum);
}


#ifdef CPU_ONLY
STUB_GPU_FORWARD(ConsNNPairLayer, Forward);
#endif

INSTANTIATE_CLASS(ConsNNPairLayer);

}  // namespace caffe
