#include <algorithm>
#include <vector>

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

namespace caffe {

template <typename Dtype>
void UnifyLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
	CHECK_EQ(bottom[0]->height(), 1);
	CHECK_EQ(bottom[0]->width(), 1);
  (*top)[0]->Reshape(bottom[0]->num(), bottom[0]->channels(),
      bottom[0]->height(), bottom[0]->width());
  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);
      temp_vec_.Reshape(bottom[0]->num(), 1, 1, 1);
      vec1.Reshape(bottom[0]->channels(), 1, 1, 1);
      vec2.Reshape(bottom[0]->channels(), 1, 1, 1);
  Dtype* multiplier_data = summer_vec_.mutable_cpu_data();
  caffe_set(summer_vec_.count(), Dtype(1), multiplier_data);
}

template <typename Dtype>
void UnifyLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
  const Dtype* bottom_data = bottom[0]->cpu_data();
  Dtype* top_data = (*top)[0]->mutable_cpu_data();
  int count=bottom[0]->count();
  const Dtype scale = this->layer_param_.unify_param().scale();
  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_powx(temp_vec_.count(), temp_vec_.cpu_data(), Dtype(0.5),
  			temp_vec_.mutable_cpu_data());
  	  int num=bottom[0]->num();
  	  int dim=bottom[0]->count()/num;
  	caffe_copy(count, bottom_data, top_data);
  	//LOG(INFO) << "...";
  	  for (int i=0;i<num;i++) {
  		//caffe_div(dim, bottom_data+bottom[0]->offset(i), temp_vec_.cpu_data(), top_data+(*top)[0]->offset(i));
  		caffe_scal(dim, scale/(temp_vec_.cpu_data()[i]), top_data+(*top)[0]->offset(i));
  	  }
}

template <typename Dtype>
void UnifyLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down,
    vector<Blob<Dtype>*>* bottom) {
  const Dtype* top_diff = top[0]->cpu_diff();

  const Dtype* top_data = top[0]->cpu_data();
  const Dtype* bottom_data = (*bottom)[0]->cpu_data();
  Dtype* bottom_diff = (*bottom)[0]->mutable_cpu_diff();
  int count=(*bottom)[0]->count();
  int num=(*bottom)[0]->num();
  int dim=(*bottom)[0]->count()/num;
  const Dtype scale = this->layer_param_.unify_param().scale();
  //caffe_copy(count, top_diff, bottom_diff);
  //for (int i=0;i<num;i++) {
  	//caffe_div(dim, top_diff+top[0]->offset(i), temp_vec_.cpu_data(), bottom_diff+(*bottom)[0]->offset(i));
  	//caffe_scal(dim, scale/(temp_vec_.cpu_data()[i]), bottom_diff+(*bottom)[0]->offset(i));
  //}
  Dtype* vec1data=vec1.mutable_cpu_data();
  Dtype* vec2data=vec2.mutable_cpu_data();
 for (int i=0;i<num;i++) {
	  caffe_copy(dim, top_diff+top[0]->offset(i), vec1data);
	  caffe_scal(dim, scale/(temp_vec_.cpu_data()[i]), vec1data);
	  caffe_copy(dim, vec1data, vec2data);
	  Dtype tmpdot=caffe_cpu_dot(dim,top_data+top[0]->offset(i),vec2data);
	  caffe_copy(dim, top_data+top[0]->offset(i), vec2data);
	  caffe_scal(dim, tmpdot, vec2data);
	  caffe_sub(dim,vec1data,vec2data, bottom_diff+(*bottom)[0]->offset(i));
  }
 /*const Dtype* top_cpu_diff = top[0]->cpu_diff();
 const Dtype* bottom_cpu_diff = (*bottom)[0]->cpu_diff();

 LOG(INFO) << this->layer_param_.name() << " " << "bottom diff:";
 for (int i=0;i<10;i++)
 {
 	LOG(INFO) << bottom_cpu_diff[i];
 }
 LOG(INFO) << this->layer_param_.name() << " " << "top diff:";
   for (int i=0;i<10;i++)
   {
   	LOG(INFO) << top_diff[i];
   }*/
 /*std::ofstream file;
 	    file.open("./record_unify.txt", ios::out|ios::app);
 	    //file << "bottom" << "\n";
 	    for (int i = 0; i < (*bottom)[0]->num(); ++i)
 	      {
 	    	  file << i << "\t";
 	    	  Dtype * bottomsample=(*bottom)[0]->offset(i)+(*bottom)[0]->mutable_cpu_data();
 	    	  for (int j=0; j<dim; ++j)
 	    	  {
 	    		  file << bottomsample[j] << "\t";
 	    	  }
 	    	  file << "\n";
 	      }
 	     //file << "top" << "\n";
 	    for (int i = 0; i < top[0]->num(); ++i)
 	     	      {
 	     	    	  file << i << "\t";
 	     	    	  Dtype * topsample=top[0]->offset(i)+top[0]->mutable_cpu_data();
 	     	    	  for (int j=0; j<dim; ++j)
 	     	    	  {
 	     	    		  file << topsample[j] << "\t";
 	     	    	  }
 	     	    	  file << "\n";
 	     	      }
 	   //file << "bottom_diff" << "\n";
 	    	      for (int i = 0; i < (*bottom)[0]->num(); ++i)
 	    	      {
 	    	    	  file << i << "\t";
 	    	    	  Dtype * bottomsample=(*bottom)[0]->offset(i)+bottom_diff;
 	    	    	  for (int j=0; j<dim; ++j)
 	    	    	  {
 	    	    		  file << bottomsample[j] << "\t";
 	    	    	  }
 	    	    	  file << "\n";
 	    	      }
 	    	     //file << "top_diff" << "\n";
 	    	      	    for (int i = 0; i < top[0]->num(); ++i)
 	    	      	     	      {
 	    	      	     	    	  file << i << "\t";
 	    	      	     	    	  Dtype * topsample=top[0]->offset(i)+top[0]->mutable_cpu_diff();
 	    	      	     	    	  for (int j=0; j<dim; ++j)
 	    	      	     	    	  {
 	    	      	     	    		  file << topsample[j] << "\t";
 	    	      	     	    	  }
 	    	      	     	    	  file << "\n";
 	    	      	     	      }
 	    	      	    file << "bottom_norm" << "\n";
 	    	      	 for (int i = 0; i < (*bottom)[0]->num(); ++i)
 	    	      	 {
 	    	      		file << temp_vec_.cpu_data()[i] << "\t";
 	    	      	 }
 	    	      	 file << "\n";
 	      file.close();*/
}


#ifdef CPU_ONLY
STUB_GPU(UnifyLayer);
#endif

INSTANTIATE_CLASS(UnifyLayer);


}  // namespace caffe
