#include <algorithm>
#include <vector>

#include "caffe/layer.hpp"
#include "caffe/loss_layers.hpp"
#include "caffe/util/io.hpp"
#include "caffe/util/math_functions.hpp"

#include "caffe/filler.hpp"

namespace caffe {

template <typename Dtype>
void GeneralizedSimilarityLossLayer<Dtype>::LayerSetUp(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  LossLayer<Dtype>::LayerSetUp(bottom, top);
  CHECK_EQ(bottom[0]->channels(), bottom[1]->channels());
  CHECK_EQ(bottom[2]->channels(), bottom[3]->channels());
  CHECK_EQ(bottom[4]->channels(), 1);
  CHECK_EQ(bottom[5]->channels(), 1);
  CHECK_EQ(bottom[0]->height(), 1);
  CHECK_EQ(bottom[0]->width(), 1);
  CHECK_EQ(bottom[1]->height(), 1);
  CHECK_EQ(bottom[1]->width(), 1);
  CHECK_EQ(bottom[2]->height(), 1);
  CHECK_EQ(bottom[2]->width(), 1);
  CHECK_EQ(bottom[3]->height(), 1);
  CHECK_EQ(bottom[3]->width(), 1);
  CHECK_EQ(bottom[4]->height(), 1);
  CHECK_EQ(bottom[4]->width(), 1);
  CHECK_EQ(bottom[5]->height(), 1);
  CHECK_EQ(bottom[5]->width(), 1);
  CHECK_EQ(bottom[6]->channels(), 1);
  CHECK_EQ(bottom[6]->height(), 1);
  CHECK_EQ(bottom[6]->width(), 1);
  bool bias_term_= this->layer_param_.generalized_similarity_loss_param().bias_term();
  if (this->blobs_.size() > 0) {
      LOG(INFO) << "Skipping parameter initialization";
    } else {
      if (bias_term_) {
        this->blobs_.resize(1);
        this->blobs_[0].reset(new Blob<Dtype>(1, 1, 1, 1));
              shared_ptr<Filler<Dtype> > bias_filler(GetFiller<Dtype>(
                  this->layer_param_.generalized_similarity_loss_param().bias_filler()));
              bias_filler->Fill(this->blobs_[0].get());
      }
    }
}

template <typename Dtype>
void GeneralizedSimilarityLossLayer<Dtype>::Reshape(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
	LossLayer<Dtype>::Reshape(bottom, top);
	//LOG(INFO) << "layersetup:" << bottom[0]->num() << "," << bottom[0]->channels() << "," << bottom[0]->height() << "," << bottom[0]->width();
	  data_squ.Reshape(bottom[0]->num(), bottom[0]->channels(), 1, 1);
	  // vector of ones used to sum along channels
	  summer_vec_.Reshape(bottom[0]->channels(), 1, 1, 1);
	  XA_norm.Reshape(bottom[0]->num(), 1, 1, 1);
	  YB_norm.Reshape(bottom[0]->num(), 1, 1, 1);
	  data23_squ.Reshape(bottom[2]->num(), bottom[2]->channels(), 1, 1);
	  for (int i = 0; i < bottom[0]->channels(); ++i)
	    summer_vec_.mutable_cpu_data()[i] = Dtype(1);
	  summer_vec_2.Reshape(bottom[2]->channels(), 1, 1, 1);
	  for (int i = 0; i < bottom[2]->channels(); ++i)
	  	    summer_vec_2.mutable_cpu_data()[i] = Dtype(1);
	  data_XCYC.Reshape(bottom[2]->num(), 1, 1, 1);
	  similarity.Reshape(bottom[0]->num(), 1, 1, 1);
}

template <typename Dtype>
void GeneralizedSimilarityLossLayer<Dtype>::Forward_cpu(
    const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
	int count1 = bottom[0]->count();
	int count2 = bottom[2]->count();
	const Dtype* bottom_data0=bottom[0]->cpu_data();
	const Dtype* bottom_data1=bottom[1]->cpu_data();
	const Dtype* bottom_data2=bottom[2]->cpu_data();
	const Dtype* bottom_data3=bottom[3]->cpu_data();
	const Dtype* bottom_data4=bottom[4]->cpu_data();
	const Dtype* bottom_data5=bottom[5]->cpu_data();
	caffe_mul(
	    count1,
		bottom_data0,
		bottom_data0,
		data_squ.mutable_cpu_data());
	caffe_cpu_gemv(
	      CblasNoTrans,
	      bottom[0]->num(),
	      bottom[0]->channels(),
	      Dtype(1.0),
		  data_squ.cpu_data(),  // a_i*a_i
	      summer_vec_.cpu_data(),
	      Dtype(0.0),
		  XA_norm.mutable_cpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1
	caffe_mul(
		    count1,
			bottom_data1,
			bottom_data1,
			data_squ.mutable_cpu_data());
		caffe_cpu_gemv(
		      CblasNoTrans,
		      bottom[0]->num(),
		      bottom[0]->channels(),
		      Dtype(1.0),
			  data_squ.cpu_data(),  // a_i*a_i
		      summer_vec_.cpu_data(),
		      Dtype(0.0),
			  YB_norm.mutable_cpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1
	caffe_mul(
		    count2,
			bottom_data2,
			bottom_data3,
			data23_squ.mutable_cpu_data());
	caffe_cpu_gemv(
		      CblasNoTrans,
		      bottom[2]->num(),
		      bottom[2]->channels(),
		      Dtype(1.0),
			  data23_squ.cpu_data(),  // a_i*a_i
		      summer_vec_2.cpu_data(),
		      Dtype(0.0),
			  data_XCYC.mutable_cpu_data());  // \Sum (a_i*a_i)=||a||^2  n1x1



	caffe_cpu_axpby(bottom[0]->num(), Dtype(1.0), XA_norm.cpu_data(), Dtype(0.0), similarity.mutable_cpu_data());
	caffe_cpu_axpby(bottom[0]->num(), Dtype(1.0), YB_norm.cpu_data(), Dtype(1.0), similarity.mutable_cpu_data());
	caffe_cpu_axpby(bottom[0]->num(), Dtype(-2.0), data_XCYC.cpu_data(), Dtype(1.0), similarity.mutable_cpu_data());
	caffe_cpu_axpby(bottom[0]->num(), Dtype(2.0), bottom_data4, Dtype(1.0), similarity.mutable_cpu_data());
	caffe_cpu_axpby(bottom[0]->num(), Dtype(2.0), bottom_data5, Dtype(1.0), similarity.mutable_cpu_data());

  Dtype margin = this->layer_param_.generalized_similarity_loss_param().margin();
  Dtype loss(0.0);
  Dtype averagedist=0;
  Dtype distthres=this->blobs_[0]->cpu_data()[0];
  for (int i = 0; i < bottom[0]->num(); ++i) {
    averagedist+=similarity.cpu_data()[i];
    if (static_cast<int>(bottom[6]->cpu_data()[i])) {  // similar pairs
      loss += std::max(similarity.cpu_data()[i]-distthres+margin, Dtype(0.0));
    } else {  // dissimilar pairs
      loss += std::max(margin+distthres-similarity.cpu_data()[i], Dtype(0.0));
    }
  }
  //LOG(INFO) << "average distance:" << averagedist/bottom[0]->num();
  loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  (*top)[0]->mutable_cpu_data()[0] = loss;
  //LOG(INFO) << "contrastive_loss:" << loss;
  if (this->layer_param_.has_recording()) {
    if (this->layer_param_.recording()) {
      std::ofstream file;
      file.open("./loss_GSL.txt", ios::out|ios::app);
        for (int i = 0; i < bottom[0]->num(); ++i)
        {
      	  file << i << "\t";
      	  file << similarity.cpu_data()[i] << "\n";
        }
        file.close();
    }
  }
}

template <typename Dtype>
void GeneralizedSimilarityLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
  Dtype margin = this->layer_param_.generalized_similarity_loss_param().margin();
  Dtype distthres=this->blobs_[0]->cpu_data()[0];
  Dtype* bottom_diff0=(*bottom)[0]->mutable_cpu_diff();
  Dtype* bottom_diff1=(*bottom)[1]->mutable_cpu_diff();
  Dtype* bottom_diff2=(*bottom)[2]->mutable_cpu_diff();
  Dtype* bottom_diff3=(*bottom)[3]->mutable_cpu_diff();
  Dtype* bottom_diff4=(*bottom)[4]->mutable_cpu_diff();
  Dtype* bottom_diff5=(*bottom)[5]->mutable_cpu_diff();
  const Dtype* bottomdata0=(*bottom)[0]->cpu_data();
  const Dtype* bottomdata1=(*bottom)[1]->cpu_data();
  const Dtype* bottomdata2=(*bottom)[2]->cpu_data();
  const Dtype* bottomdata3=(*bottom)[3]->cpu_data();
  const Dtype* bottomdata4=(*bottom)[4]->cpu_data();
  const Dtype* bottomdata5=(*bottom)[5]->cpu_data();
  int dim01=(*bottom)[0]->count()/(*bottom)[0]->num();
  int dim23=(*bottom)[2]->count()/(*bottom)[2]->num();

  caffe_set((*bottom)[0]->count(), Dtype(0.0), bottom_diff0);
  caffe_set((*bottom)[1]->count(), Dtype(0.0), bottom_diff1);
  caffe_set((*bottom)[2]->count(), Dtype(0.0), bottom_diff2);
  caffe_set((*bottom)[3]->count(), Dtype(0.0), bottom_diff3);
  caffe_set((*bottom)[4]->count(), Dtype(0.0), bottom_diff4);
  caffe_set((*bottom)[5]->count(), Dtype(0.0), bottom_diff5);

  Dtype* dt_diff=this->blobs_[0]->mutable_cpu_diff();
  dt_diff[0]=Dtype(0.0);

  int num=(*bottom)[0]->num();

  for (int j=0;j<num;j++) {
	  if (static_cast<int>((*bottom)[6]->cpu_data()[j])) {//similarity
		  if ((similarity.cpu_data()[j]-distthres+margin) > Dtype(0.0)) {
			  const Dtype alpha = top[0]->cpu_diff()[0] / static_cast<Dtype>((*bottom)[0]->num());
			  caffe_cpu_axpby(dim01, alpha, bottomdata0+(*bottom)[0]->offset(j), Dtype(1.0), bottom_diff0+(*bottom)[0]->offset(j));
			  caffe_cpu_axpby(dim01, alpha, bottomdata1+(*bottom)[1]->offset(j), Dtype(1.0), bottom_diff1+(*bottom)[1]->offset(j));
			  caffe_cpu_axpby(dim23, -alpha, bottomdata3+(*bottom)[3]->offset(j), Dtype(1.0), bottom_diff2+(*bottom)[2]->offset(j));
			  caffe_cpu_axpby(dim23, -alpha, bottomdata2+(*bottom)[2]->offset(j), Dtype(1.0), bottom_diff3+(*bottom)[3]->offset(j));
			  bottom_diff4[j]=bottom_diff4[j]+alpha;
			  bottom_diff5[j]=bottom_diff5[j]+alpha;
			  dt_diff[0]=dt_diff[0]-Dtype(0.5) * alpha;
		  }
	  } else {//dissimilarity
		  if ((distthres+margin-similarity.cpu_data()[j]) > Dtype(0.0)) {
			  const Dtype alpha = top[0]->cpu_diff()[0] / static_cast<Dtype>((*bottom)[0]->num());
			  caffe_cpu_axpby(dim01, -alpha, bottomdata0+(*bottom)[0]->offset(j), Dtype(1.0), bottom_diff0+(*bottom)[0]->offset(j));
			  caffe_cpu_axpby(dim01, -alpha, bottomdata1+(*bottom)[1]->offset(j), Dtype(1.0), bottom_diff1+(*bottom)[1]->offset(j));
			  caffe_cpu_axpby(dim23, alpha, bottomdata3+(*bottom)[3]->offset(j), Dtype(1.0), bottom_diff2+(*bottom)[2]->offset(j));
			  caffe_cpu_axpby(dim23, alpha, bottomdata2+(*bottom)[2]->offset(j), Dtype(1.0), bottom_diff3+(*bottom)[3]->offset(j));
			  bottom_diff4[j]=bottom_diff4[j]-alpha;
			  bottom_diff5[j]=bottom_diff5[j]-alpha;
			  dt_diff[0]=dt_diff[0]+Dtype(0.5) * alpha;
		  }
	  }
  }


}

#ifdef CPU_ONLY
STUB_GPU(GeneralizedSimilarityLossLayer);
#endif

INSTANTIATE_CLASS(GeneralizedSimilarityLossLayer);

}  // namespace caffe
