// Copyright 2014 BVLC and contributors.

#include <algorithm>
#include <cmath>
#include <cfloat>
#include <vector>

#include "caffe/layer.hpp"
#include "caffe/vision_layers.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/io.hpp"
#include "caffe/filler.hpp"

using std::max;

namespace caffe {

template <typename Dtype>
void EuclideanCombLayer<Dtype>::LayerSetUp(
  const vector<Blob<Dtype>*>& bottom, vector<Blob<Dtype>*>* top) {
  CHECK_EQ(bottom[0]->channels(), bottom[1]->channels());
  CHECK_EQ(bottom[0]->height(), bottom[1]->height());
  CHECK_EQ(bottom[0]->width(), bottom[1]->width());
  M_ = bottom[0]->num();
  K_ = bottom[0]->count() / bottom[0]->num();
  if (this->blobs_.size() > 0) {
      LOG(INFO) << "Skipping parameter initialization";
  } else {
      this->blobs_.resize(1);
    // Intialize the bias
      this->blobs_[0].reset(new Blob<Dtype>(1, 1, 1, 1));
      shared_ptr<Filler<Dtype> > bias_filler(GetFiller<Dtype>(
          this->layer_param_.euclidean_threshold_param().bias_filler()));
      bias_filler->Fill(this->blobs_[0].get());
      Dtype* bias_data = this->blobs_[0]->mutable_cpu_data();
      bias_data[0]=1e+6;
  }  // parameter initialization
}

template <typename Dtype>
void EuclideanCombLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      vector<Blob<Dtype>*>* top) {
	  (*top)[0]->Reshape(M_*M_, 1, 1, 1);
	  (*top)[1]->Reshape(M_*M_, 1, 1, 1);
	  b1norm.Reshape(bottom[0]->num(), 1, 1, 1);
	  b2norm.Reshape(bottom[0]->num(), 1, 1, 1);
	  diff_.Reshape(1, bottom[0]->channels(),
	      bottom[0]->height(), bottom[0]->width());
}

template <typename Dtype>
void EuclideanCombLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    vector<Blob<Dtype>*>* top) {
	int count = bottom[0]->count();
	int channels = bottom[0]->channels();
	int height = bottom[0]->height();
	int width = bottom[0]->width();
	const Dtype* bottom_data1=bottom[0]->cpu_data();
	const Dtype* bottom_data2=bottom[1]->cpu_data();
	const Dtype* label = bottom[2]->cpu_data();
	Dtype * diff_data=diff_.mutable_cpu_data();
	/*caffe_sub(
		count,
		bottom[0]->cpu_data(),
		bottom[1]->cpu_data(),
		diff_.mutable_cpu_data());*/
	const Dtype* bias_data = this->blobs_[0]->cpu_data();
	//Dtype dot = caffe_cpu_dot(count, diff_.cpu_data(), diff_.cpu_data());
	Dtype* dist_data=(*top)[0]->mutable_cpu_data();
	Dtype * pairlabel = (*top)[1]->mutable_cpu_data();
	for (int i=0;i<M_;i++){
		for (int j=0; j<M_; j++){
			caffe_sub(K_,bottom_data1+(i*K_),bottom_data2+(j*K_), diff_data);
		dist_data[i*M_+j]=caffe_cpu_dot(K_,diff_data, diff_data);
		dist_data[i*M_+j]=dist_data[i*M_+j]-bias_data[0];
		pairlabel[i*M_+j]=(label[i]==label[j])?1:-1;
		}
	}
}

template <typename Dtype>
void EuclideanCombLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
		const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {
	Dtype* diff_data=diff_.mutable_cpu_data();
	const Dtype* top_diff = top[0]->cpu_diff();
	Dtype* bias_diff = this->blobs_[0]->mutable_cpu_diff();
	const Dtype* bottom_data1=(*bottom)[0]->cpu_data();
	const Dtype* bottom_data2=(*bottom)[1]->cpu_data();
	Dtype* bottom_diff1=(*bottom)[0]->mutable_cpu_diff();
	Dtype* bottom_diff2=(*bottom)[1]->mutable_cpu_diff();
	caffe_scal((*bottom)[0]->count(),Dtype(0),bottom_diff1);
	caffe_scal((*bottom)[1]->count(),Dtype(0),bottom_diff2);
	bias_diff[0]=Dtype(0);
	for (int i=0; i<M_;i++)
	{
		for (int j=0; j<M_; j++)
		{
			caffe_sub(K_,bottom_data1+(i*K_),bottom_data2+(j*K_), diff_data);
			caffe_cpu_axpby(K_, Dtype(2*top_diff[i*M_+j]), diff_data, Dtype(1), bottom_diff1+(i*K_));
			caffe_cpu_axpby(K_, Dtype(-2*top_diff[i*M_+j]), diff_data, Dtype(1), bottom_diff1+(i*K_));
			bias_diff[0]=bias_diff[0]-top_diff[i*M_+j];
		}
	}
}

#ifdef CPU_ONLY
STUB_GPU(EuclideanCombLayer);
#endif

INSTANTIATE_CLASS(EuclideanCombLayer);

}  // namespace caffe
