#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 KNNAccuracyLayer<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();
  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();
  caffe_gpu_mul(
    probe_count,
    probe_data,
    probe_data,
    probe_squ.mutable_gpu_data());
  caffe_gpu_mul(
    gallery_count,
    gallery_data,
    gallery_data,
    gallery_squ.mutable_gpu_data());
  caffe_gpu_gemv(
      CblasNoTrans,
      bottom[0]->num(),
      bottom[0]->channels(),
      Dtype(1.0),
      probe_squ.gpu_data(),  // a_i*a_i
      summer_vec_.gpu_data(),
      Dtype(0.0),
      probe_norm.mutable_gpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1
  caffe_gpu_gemv(
      CblasNoTrans,
      bottom[1]->num(),
      bottom[1]->channels(),
      Dtype(1.0),
      gallery_squ.gpu_data(),  // b_i*b_i
      summer_vec_.gpu_data(),
      Dtype(0.0),
      gallery_norm.mutable_gpu_data());  // \Sum (b_i*b_i)=||b||^2  n2x1
  caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[0]->num(), bottom[1]->num(), 1,
      Dtype(1.0),
      probe_norm.gpu_data(), summer_vec_gallery.gpu_data(),
      Dtype(0),
      dist.mutable_gpu_data());
  caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[0]->num(), bottom[1]->num(), 1,
      Dtype(1.0),
      summer_vec_probe.gpu_data(), gallery_norm.gpu_data(),
      Dtype(1.0),
      dist.mutable_gpu_data());
  caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[0]->num(), bottom[1]->num(), bottom[0]->channels(),
      Dtype(-2.0),
      probe_data, gallery_data,
      Dtype(1.0),
      dist.mutable_gpu_data());
  Dtype* dist_data=dist.mutable_cpu_data();
  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++) {
      Dtype* dist_ij=dist_data+dist.offset(i,j);
          if (dist_ij[0]<mindist) {
              mindist=dist_ij[0];
              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();
  
  string lossfile(this->layer_param_.name());
  lossfile="./"+lossfile+".txt";
  string lossfile1="./"+lossfile+"_1.txt";
  string lossfile2="./"+lossfile+"_2.txt";
  if (this->layer_param_.has_recording()) {
  if (this->layer_param_.recording()) {
  //LOG(INFO) << "cl recording";
    std::ofstream file;
    file.open(lossfile.data(), 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();
      std::ofstream file1,file2;
      Dtype* probe_data_cpu = bottom[0]->mutable_cpu_data();
      Dtype* gallery_data_cpu = bottom[1]->mutable_cpu_data();
      file1.open(lossfile1.data(), ios::out|ios::app);
      //file2.open(lossfile2.data(), ios::out|ios::app);
      for (int i = 0; i < bottom[0]->num(); ++i)
      {
        Dtype* probe_data_i=probe_data_cpu+bottom[0]->offset(i);
        for (int j=0;j<bottom[0]->channels(); ++j)
        {
          file1<<probe_data_i[j] << "\t";
        }
        file1 << "\n";
      }
      /*for (int i = 0; i < bottom[1]->num(); ++i)
      {
        Dtype* gallery_data_i=gallery_data_cpu+bottom[1]->offset(i);
        for (int j=0;j<bottom[1]->channels(); ++j)
        {
          file2<<gallery_data_i[j] << "\t";
        }
        file2 << "\n";
      }*/
      file1.close();
  }
  }
}



INSTANTIATE_CLASS(KNNAccuracyLayer);

}  // namespace caffe
