#include <algorithm>
#include <vector>

#include "caffe/common_layers.hpp"
#include "caffe/layer.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

template <typename Dtype>
__global__ void ULForward(const int count, const int dim,
    const Dtype* x, const Dtype* dist, Dtype* y, Dtype scale) {
  CUDA_KERNEL_LOOP(i, count) {
    int n = i / dim;  // the num index, to access y and dist_sq
    y[i] = x[i] / dist[n] * scale;
  }
}

template <typename Dtype>
void UnifyLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
    const Dtype* bottom_data = bottom[0]->gpu_data();
  Dtype* top_data = (*top)[0]->mutable_gpu_data();
  int count=bottom[0]->count();
  const Dtype scale = this->layer_param_.unify_param().scale();
  caffe_gpu_mul(
  	      count,
  	      bottom[0]->gpu_data(),  // a
  	      bottom[0]->gpu_data(),  // a
  	      temp_mat_.mutable_gpu_data());  // a_i*a_i
  	    caffe_gpu_gemv(
  	      CblasNoTrans,
  	      bottom[0]->num(), bottom[0]->channels(),
  	      Dtype(1.0),
  	      temp_mat_.gpu_data(),  // a_i*a_i
  	      summer_vec_.gpu_data(),
  	      Dtype(0.0),
  	      temp_vec_.mutable_gpu_data());  // \Sum a_i*a_i
  	  caffe_gpu_powx(temp_vec_.count(), temp_vec_.gpu_data(), Dtype(0.5),
  			temp_vec_.mutable_gpu_data());
  	  int num=bottom[0]->num();
  	  //Dtype sumvec=0;
  	  //for (int i=0;i<num;i++) {
  	   // sumvec=sumvec+temp_vec_.cpu_data()[i];
  	 // }
  	  //LOG(INFO) << "sumvec:" << sumvec/num;
  	  int dim=count/num;
  	  // NOLINT_NEXT_LINE(whitespace/operators)
      ULForward<Dtype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
          count, dim, 
          bottom_data,  
          temp_vec_.gpu_data(),  
          top_data,
          scale);
      CUDA_POST_KERNEL_CHECK;
      /*LOG(INFO) << this->layer_param_.name() << " bottom:";
      const Dtype* bottom_cpu_data = bottom[0]->cpu_data();
      const Dtype* top_cpu_data = (*top)[0]->cpu_data();
      for (int i=0;i<10;i++) {
        LOG(INFO) << bottom_cpu_data[i];
      }
      LOG(INFO) << this->layer_param_.name() << " top:";
      for (int i=0;i<10;i++) {
        LOG(INFO) << top_cpu_data[i];
      }*/
}

template <typename Dtype>
void UnifyLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down,
    vector<Blob<Dtype>*>* bottom) {
    Backward_cpu(top,propagate_down,bottom);
      /* const Dtype* top_diff = top[0]->gpu_diff();
  const Dtype* top_data = top[0]->gpu_data();
  const Dtype* bottom_data = (*bottom)[0]->gpu_data();
  Dtype* bottom_diff = (*bottom)[0]->mutable_gpu_diff();
  const Dtype scale = this->layer_param_.unify_param().scale();
  int num=(*bottom)[0]->num();
  	  int count=(*bottom)[0]->count();
  	  int dim=count/num;
  // NOLINT_NEXT_LINE(whitespace/operators)
   ULForward<Dtype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
          count, dim, 
          top_diff,  
          temp_vec_.gpu_data(),  
          bottom_diff,
          scale);
      CUDA_POST_KERNEL_CHECK;*/
}


INSTANTIATE_CLASS(UnifyLayer);


}  // namespace caffe
