#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 GSKNNAccuracyLayer<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* XC = bottom[2]->gpu_data();
  const Dtype* YC = bottom[3]->gpu_data();
  const Dtype* XD = bottom[4]->gpu_data();
  const Dtype* YE = bottom[5]->gpu_data();
  const Dtype* probe_label = bottom[6]->cpu_data();
  const Dtype* gallery_label = bottom[7]->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),
      XC, YC,
      Dtype(1.0),
      dist.mutable_gpu_data());
  caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[4]->num(), bottom[5]->num(), 1,
      Dtype(2.0),
      XD, summer_vec_gallery.gpu_data(),
      Dtype(1.0),
      dist.mutable_gpu_data());
  caffe_gpu_gemm(
      CblasNoTrans, CblasTrans,
      bottom[4]->num(), bottom[5]->num(), 1,
      Dtype(2.0),
      summer_vec_probe.gpu_data(), YE,
      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();
  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(GSKNNAccuracyLayer);

}  // namespace caffe
