#include <vector>

#include <cmath>
#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 EuclideanLossLayer<Dtype>::Reshape(
  const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
  LossLayer<Dtype>::Reshape(bottom, top);
  CHECK_EQ(bottom[0]->count(1), bottom[1]->count(1))
      << "Inputs must have the same dimension.";
  diff_.ReshapeLike(*bottom[0]);
}

template <typename Dtype>
void EuclideanLossLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {

  //if(this->phase_ == TEST) {
  //  int max_show = 100;
  //  int show = bottom[0]->count() > max_show ? max_show : bottom[0]->count();
  //  for (int i = 0; i < show; ++i) {
  //    LOG(INFO) << bottom[0]->cpu_data()[i] << "\t" << bottom[1]->cpu_data()[i];
  //  }
  //}

  //Dtype accuracy = 0;
  //for (int i = 0; i < bottom[0]->count(); ++i) {
  //  int pred = round(bottom[0]->cpu_data()[i]);
  //  int truth = round(bottom[1]->cpu_data()[i]);
  //  if (pred == truth) {
  //    accuracy += 1;
  //  }
  //}
  //LOG(INFO) << "Accuracy\t" << (accuracy / bottom[0]->count())
  //    << "\t" << accuracy << "\t" << bottom[0]->count();

  int count = bottom[0]->count();
  caffe_sub(
      count,
      bottom[0]->cpu_data(),
      bottom[1]->cpu_data(),
      diff_.mutable_cpu_data());
  Dtype dot = caffe_cpu_dot(count, diff_.cpu_data(), diff_.cpu_data());
  Dtype loss = dot / bottom[0]->num() / Dtype(2);
  top[0]->mutable_cpu_data()[0] = loss;
 
  if(isinf(loss) || isnan(loss)) {
    LOG(ERROR) << dot << "\t" << bottom[0]->num();
  }
}

template <typename Dtype>
void EuclideanLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  for (int i = 0; i < 2; ++i) {
    if (propagate_down[i]) {
      const Dtype sign = (i == 0) ? 1 : -1;
      const Dtype alpha = sign * top[0]->cpu_diff()[0] / bottom[i]->num();
      caffe_cpu_axpby(
          bottom[i]->count(),              // count
          alpha,                              // alpha
          diff_.cpu_data(),                   // a
          Dtype(0),                           // beta
          bottom[i]->mutable_cpu_diff());  // b
    }
  }
}

#ifdef CPU_ONLY
STUB_GPU(EuclideanLossLayer);
#endif

INSTANTIATE_CLASS(EuclideanLossLayer);
REGISTER_LAYER_CLASS(EuclideanLoss);

}  // namespace caffe
