#include <algorithm>
#include <vector>

#include "caffe/layers/contrastive_accuracy_layer.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

template <typename Dtype>
void ContrastiveAccuracyLayer<Dtype>::Forward_gpu(
    const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
  const int count = bottom[0]->count();
  caffe_gpu_sub(
      count,
      bottom[0]->gpu_data(),  // a
      bottom[1]->gpu_data(),  // b
      diff_.mutable_gpu_data());  // a_i-b_i
  caffe_gpu_powx(
      count,
      diff_.mutable_gpu_data(),  // a_i-b_i
      Dtype(2),
      diff_sq_.mutable_gpu_data());  // (a_i-b_i)^2
  caffe_gpu_gemv(
      CblasNoTrans,
      bottom[0]->num(),
      bottom[0]->channels(),
      Dtype(1.0),
      diff_sq_.gpu_data(),  // (a_i-b_i)^2
      summer_vec_.gpu_data(),
      Dtype(0.0),
      dist_sq_.mutable_gpu_data());  // \Sum (a_i-b_i)^2
  Dtype margin = this->layer_param_.contrastive_accuracy_param().margin();
  bool legacy_version =
      this->layer_param_.contrastive_accuracy_param().legacy_version();
  // Dtype loss(0.0);
  int pos_cnt = 0;
  int neg_cnt = 0;
  int pos_right = 0;
  int neg_right = 0;
  float eps = 0.0001;

  for (int i = 0; i < bottom[0]->num(); ++i) {
    if (static_cast<int>(bottom[2]->cpu_data()[i])) {  // similar pairs
      // loss += dist_sq_.cpu_data()[i];
      // handle the postive pair
      pos_cnt += 1;
      if (dist_sq_.cpu_data()[i] < margin) pos_right += 1;

    } else {  // dissimilar pairs
      if (legacy_version) {
        // loss += std::max(margin - dist_sq_.cpu_data()[i], Dtype(0.0));
        neg_cnt += 1;
        if( std::max(margin - dist_sq_.cpu_data()[i], Dtype(0.0)) == 0)
        {
            neg_right += 1;
        }
      } else {
        Dtype dist = std::max(margin - sqrt(dist_sq_.cpu_data()[i]),
                              Dtype(0.0));
        // loss += dist*dist;
        neg_cnt += 1;
        if (dist == 0)
        {
            neg_right += 1;
        }
      }
    }
  }
  // loss = loss / static_cast<Dtype>(bottom[0]->num()) / Dtype(2);
  float pos_accuracy = pos_right/(pos_cnt + eps);
  float neg_accuracy = neg_right/(neg_cnt + eps);
  float accuracy = 0.5*(pos_accuracy + neg_accuracy);
  top[0]->mutable_cpu_data()[0] = accuracy;
  top[0]->mutable_cpu_data()[1] = pos_accuracy;
  top[0]->mutable_cpu_data()[2] = neg_accuracy;
}

template <typename Dtype>
__global__ void CLLBackward(const int count, const int channels,
    const Dtype margin, const bool legacy_version, const Dtype alpha,
    const Dtype* y, const Dtype* diff, const Dtype* dist_sq,
    Dtype *bottom_diff) {
  CUDA_KERNEL_LOOP(i, count) {
    int n = i / channels;  // the num index, to access y and dist_sq
    if (static_cast<int>(y[n])) {  // similar pairs
      bottom_diff[i] = alpha * diff[i];
    } else {  // dissimilar pairs
      Dtype mdist(0.0);
      Dtype beta(0.0);
      if (legacy_version) {
        mdist = (margin - dist_sq[n]);
        beta = -alpha;
      } else {
        Dtype dist = sqrt(dist_sq[n]);
        mdist = (margin - dist);
        beta = -alpha * mdist / (dist + Dtype(1e-4)) * diff[i];
      }
      if (mdist > 0.0) {
        bottom_diff[i] = beta;
      } else {
        bottom_diff[i] = 0;
      }
    }
  }
}

template <typename Dtype>
void ContrastiveAccuracyLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom)
{
    for (int i = 0; i < propagate_down.size(); ++i) {
      if (propagate_down[i]) { NOT_IMPLEMENTED; }
    }
}
INSTANTIATE_LAYER_GPU_FUNCS(ContrastiveAccuracyLayer);

}  // namespace caffe
