#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 GSV2KNNAccuracyLayer<Dtype>::Forward_gpu(
    const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  const int probe_count = bottom[0]->count();
  const int gallery_count = bottom[1]->count();
  int size=bottom[0]->count()/bottom[0]->num();
  const Dtype* probe_data = bottom[0]->gpu_data();
  const Dtype* gallery_data = bottom[1]->gpu_data();

  const Dtype* probe_label = bottom[2]->cpu_data();
  const Dtype* gallery_label = bottom[3]->cpu_data();
  /*for (int i=0;i<20;i++)
  {
    LOG(INFO) << probe_label[i] << "," << gallery_label[i];
  }*/
  Blob<Dtype> dist;
  dist.Reshape(bottom[0]->num(), bottom[1]->num(), 1, 1);
  Dtype* dist_data=dist.mutable_cpu_data();
  for (int i=0;i<bottom[0]->num();i++) {
    Dtype dist1,dist2,dist3;
    Dtype* dist_i=dist_data+dist.offset(i);
    caffe_gpu_dot((int)(size/2), probe_data+bottom[0]->offset(i,0), probe_data+bottom[0]->offset(i,0), &dist1);
    for (int j=0;j<bottom[1]->num();j++) {
      caffe_gpu_dot((int)(size/2), gallery_data+bottom[1]->offset(j,0), gallery_data+bottom[1]->offset(j,0), &dist2);
      caffe_gpu_dot((int)(size/2), probe_data+bottom[0]->offset(i,(int)(size/2)), gallery_data+bottom[1]->offset(j,(int)(size/2)), &dist3);
      dist_i[j]=0.5*dist1+0.5*dist2-dist3;
    }
  }
  
  Dtype correctnum=0;
  for (int i=0;i<bottom[0]->num();i++) {
      Dtype* dist_i=dist_data+dist.offset(i);
      Dtype mindist=dist_i[0];
      int minindex=0;
      for (int j=0;j<bottom[1]->num();j++) {
          if (dist_i[j]<mindist) {
              mindist=dist_i[j];
              minindex=j;
          }
      }
      if (probe_label[i]==gallery_label[minindex]) {
          correctnum=correctnum+1;
      }
  }
      LOG(INFO) << "accuracy: " << correctnum/bottom[0]->num();
  (*top)[0]->mutable_cpu_data()[0] = correctnum/bottom[0]->num();
  for (int i=0;i<10; ++i) {
    LOG(INFO) << dist.cpu_data()[i];
  }
  if (this->layer_param_.has_recording()) {
  if (this->layer_param_.recording()) {
  //LOG(INFO) << "cl recording";
    std::ofstream file;
    file.open("./loss_dist.txt", ios::out|ios::app);
      for (int i = 0; i < bottom[0]->num(); ++i)
      {
    	  Dtype* dist_i=dist_data+dist.offset(i);
    	  for (int j=0;j<bottom[1]->num();j++) {
    	      file << dist_i[j] << "\t";
    	  }
    	  file << "\n";
      }
      file.close();
  }
  }
}



INSTANTIATE_CLASS(GSV2KNNAccuracyLayer);

}  // namespace caffe
