/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once
#include "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenMatrix = framework::EigenMatrix<T, MajorType, IndexType>;

template <typename Place, typename T>
class SquaredL2DistanceKernel : public framework::OpKernel {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    auto* in0 = context.Input<Tensor>("X");
    auto* in1 = context.Input<Tensor>("Y");
    auto* out0 = context.Output<Tensor>("sub_result");
    auto* out1 = context.Output<Tensor>("Out");

    auto in0_dims = in0->dims();
    auto in1_dims = in1->dims();

    int cols = framework::product(in0_dims) / in0_dims[0];
    // reduce dimensions except the first
    auto x =
        EigenMatrix<T>::From(*in0, framework::make_ddim({in0_dims[0], cols}));
    auto y =
        EigenMatrix<T>::From(*in1, framework::make_ddim({in1_dims[0], cols}));

    out0->mutable_data<T>(context.GetPlace());
    out1->mutable_data<T>(context.GetPlace());
    auto sub_result = EigenMatrix<T>::From(*out0);
    auto z = EigenMatrix<T>::From(*out1);

    auto place = context.GetEigenDevice<Place>();
    auto x_dims = x.dimensions();
    auto y_dims = y.dimensions();
    // buffer the substraction result
    if (y_dims[0] == 1 && x_dims[0] > y_dims[0]) {
      sub_result.device(place) =
          x -
          y.broadcast(Eigen::array<int, 2>({static_cast<int>(x_dims[0]), 1}));
    } else {
      sub_result.device(place) = x - y;
    }
    auto sub_res_pow2 = sub_result * sub_result;
    // z is TensorMap, no need reshape
    z.device(place) = sub_res_pow2.sum(Eigen::array<int, 1>({1}));
  }
};

template <typename Place, typename T>
class SquaredL2DistanceGradKernel : public framework::OpKernel {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    auto* in0 = context.Input<Tensor>("sub_result");
    auto* in1 = context.Input<Tensor>(framework::GradVarName("Out"));
    auto* x_g = context.Output<Tensor>(framework::GradVarName("X"));
    auto* y_g = context.Output<Tensor>(framework::GradVarName("Y"));

    auto sub_result = EigenMatrix<T>::From(*in0);
    auto out_grad = EigenMatrix<T>::From(*in1);

    auto x_dims = x_g->dims();
    auto y_dims = y_g->dims();

    int cols = framework::product(x_dims) / x_dims[0];
    // calculate gradient
    auto grad_mat =
        2 * (out_grad.broadcast(Eigen::array<int, 2>({1, cols}))) * sub_result;

    // propagate back to input
    auto eigen_place = context.GetEigenDevice<Place>();
    if (x_g) {
      x_g->mutable_data<T>(context.GetPlace());
      // eigen matrix
      auto x_grad =
          EigenMatrix<T>::From(*x_g, framework::make_ddim({x_dims[0], cols}));
      // dimensions are same with subResult
      x_grad.device(eigen_place) = grad_mat;
    }

    if (y_g) {
      y_g->mutable_data<T>(context.GetPlace());
      auto y_grad =
          EigenMatrix<T>::From(*y_g, framework::make_ddim({y_dims[0], cols}));

      PADDLE_ENFORCE(sub_result.dimensions()[0] >= y_dims[0],
                     "First dimension of gradient must be greater or "
                     "equal than first dimension of target");

      if (sub_result.dimensions()[0] == y_dims[0]) {
        y_grad.device(eigen_place) = -1 * grad_mat;
      } else {
        auto col_sum_res = -1 * (grad_mat.sum(Eigen::array<int, 1>({0})));
        // y_grad is TensorMap, no need reshape
        y_grad.device(eigen_place) = col_sum_res;
      }
    }
  }
};

}  // namespace operators
}  // namespace paddle