// 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 ConsPairLayer<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 ConsPairLayer<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
    const int NN=10;
    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);
    //LOG(INFO) << "layersetup:" << bottom[0]->num() << "," << bottom[0]->channels() << "," << bottom[0]->height() << "," << bottom[0]->width();
    for (int i=0;i<bottom[0]->num();i++)
    {
    	if (label[i]>maxlabel) {maxlabel=label[i];}
    	if (label[i]<minlabel) {minlabel=label[i];}
    }
    Dtype* labelnum=new Dtype[maxlabel+1];
    for (int j=0;j<=maxlabel;j++)
    {
    	labelnum[j]=0;
    }
    for (int i=0;i<bottom[0]->num();i++)
    {
    	labelnum[int(label[i])]++;
    }
    pospairnum=0;
    for (int j=0;j<=maxlabel;j++)
    {
    	pospairnum+=labelnum[j]*(labelnum[j]-1)/2;
    }
    if (pospairnum==0)
    {
    	LOG(INFO) << "warning: no pair constructed!";
    	pospairnum=bottom[0]->num()/2;
    }
    negpairnum=bottom[0]->num()*NN;
    //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 ConsPairLayer<Dtype>::Forward_cpu(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;
		int dim=bottom[0]->count()/bottom[0]->num();
		Dtype* data1idxarray=(*top)[0]->mutable_cpu_data();
		Dtype* data2idxarray=(*top)[1]->mutable_cpu_data();
		int nearneighbor[NN];
		Dtype nndist[NN];
		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<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++;pairidx1++;}
	          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]=j; pairlabel[pairidx]=Dtype(0.0); pairidx++;pairidx2++;}
	      }
	    }
	    //LOG(INFO) << pairidx1 << "," << pairidx2;
	    CHECK_EQ(pairidx,pospairnum+negpairnum);
}


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

INSTANTIATE_CLASS(ConsPairLayer);

}  // namespace caffe
