#include "caffe2/operators/pow_op.h"
#include "caffe2/utils/eigen_utils.h"
#include "caffe2/utils/math.h"
// definition of NumericTypes and SameTypeAsInput is in below header file
//#include "caffe2/operators/elementwise_op.h"
#include <Eigen/Core>

namespace caffe2 {

#define EIGEN_POW(x, y) (x.pow(y))

struct EigenPowFunctor {
  template <int b_is_scalar, typename T1, typename T2, typename R>
  inline void
  Run(size_t n, const T1* a, const T2* b, T2 e, R* out, CPUContext*) {
    if (b == NULL) {
      EigenVectorArrayMap<R>(out, n) =
          EIGEN_POW((ConstEigenVectorArrayMap<T1>(a, n)), (e));
    } else {
      if (b_is_scalar) {
        EigenVectorArrayMap<R>(out, n) =
            EIGEN_POW((ConstEigenVectorArrayMap<T1>(a, n)), (b[0]));
      } else {
        EigenVectorArrayMap<R>(out, n) = EIGEN_POW(
            (ConstEigenVectorArrayMap<T1>(a, n)),
            (ConstEigenVectorArrayMap<T2>(b, n)));
      }
    }
  }
  template <typename T1, typename T2, typename R>
  void RunWithBroadcast(
      const T1* a,
      const T2* b,
      R* out,
      size_t pre,
      size_t n,
      CPUContext*) {
    EigenArrayMap<R>(out, n, pre) = EIGEN_POW(
        (ConstEigenArrayMap<T1>(a, n, pre)),
        (ConstEigenVectorArrayMap<T2>(b, n)).rowwise().replicate(pre));
    /*
    //below code only allows elementary ops, such as +, -, * and /,
    //and does not allow operations, such as pow, exp and log
    EIGEN_POW(
       (ConstEigenArrayMap<T>(a, n, pre).colwise()),
       (ConstEigenVectorArrayMap<T>(b, n)));
     */
  }
  template <typename T1, typename T2, typename R>
  void RunWithBroadcast2(
      const T1* a,
      const T2* b,
      R* out,
      size_t pre,
      size_t n,
      size_t post,
      CPUContext*) {
    for (int i = 0; i < pre; ++i) {
      EigenArrayMap<R>(out + i * n * post, post, n) = EIGEN_POW(
          (ConstEigenArrayMap<T1>(a + i * n * post, post, n)),
          (Eigen::Map<const Eigen::Array<T2, 1, Eigen::Dynamic>>(b, n))
              .colwise()
              .replicate(post));
      /*
      //below code only allows elementary ops, such as +, -, * and /,
      //and does not allow for operations, such as pow, exp and log
      EIEGN_POW(
        (ConstEigenArrayMap<T>(a + i * n * post, post, n).rowwise()),
        (Eigen::Map<const Eigen::Array<T, 1, Eigen::Dynamic>>(b, n)));
      */
    }
  }
};

REGISTER_CPU_OPERATOR(
    Pow,
    PowOp<
        TensorTypes<float> /*NumericTypes*/,
        CPUContext,
        EigenPowFunctor,
        SameTypeAsInput>)

OPERATOR_SCHEMA(Pow)
    .NumInputs(1, 2)
    .NumOutputs(1)
    .AllowInplace({{0, 0}, {1, 0}})
    .IdenticalTypeAndShapeOfInput(0)
    .SetDoc(R"DOC(
The *Pow* op takes an input data tensor $X$ and an exponent parameter *exponent*, which can be a scalar or another tensor. As output, it produces a single output data tensor $Y$, where the function $f(x) = x^{exponent}$ has been applied to $X$ elementwise.

Github Links:

- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/pow_op.h
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/pow_op.cc


<details>

<summary> <b>Example</b> </summary>

**Code**

```

workspace.ResetWorkspace()

op = core.CreateOperator(
    "Pow",
    ["X", "exponent"],
    ["Y"],
    broadcast=1
)

workspace.FeedBlob("X", np.array([1,2,3,4,5,6]).astype(np.float32))
print("X: ", workspace.FetchBlob("X"))

workspace.FeedBlob("exponent", np.array([2]).astype(np.float32))
print("exponent: ", workspace.FetchBlob("exponent"))

workspace.RunOperatorOnce(op)
print("Y: ", workspace.FetchBlob("Y"))

```

**Result**

```

X:  [1. 2. 3. 4. 5. 6.]
exponent:  [2.]
Y:  [ 1.  4.  9. 16. 25. 36.]

```

</details>


)DOC")
    .Input(0, "X", "Input data blob to be operated on.")
    .Input(1, "exponent", "Exponent blob containing the exponent(s) for calculation. Do not use if setting exponent via argument.")
    .Output(0, "Y", "Output data blob with the same shape as the input.")
    .Arg("exponent", "The exponent of the power function. Do not use if setting exponent via input.")
    .Arg("axis", "*(type: int; default: -1)*")
    .Arg("broadcast", "*(type: bool; default: False)*");

class GetPowGradient : public GradientMakerBase {
  using GradientMakerBase::GradientMakerBase;
  vector<OperatorDef> GetGradientDefs() override {
    ArgumentHelper arg_helper(def_);
    if (arg_helper.HasArgument("exponent")) { // second input is a scalar
      // function f(w,a) = w^a
      // gradient operator with respect to first input tensor
      // df/dw = a * w^(a-1) (all operations are component-wise)
      float exponent = arg_helper.GetSingleArgument<float>("exponent", 0.0);
      Argument scale_arg;
      scale_arg.set_name("scale");
      scale_arg.set_f(exponent);
      Argument pow_arg;
      pow_arg.set_name("exponent");
      if (I(0) != O(0)) {
        pow_arg.set_f(exponent - 1);
      } else {
        LOG(WARNING) << "In-place Pow gradient, possible loss of precision";
        constexpr float kEps = 1e-12f;
        CAFFE_ENFORCE(std::fabs(exponent) > kEps);
        pow_arg.set_f((exponent - 1) / exponent);
      }
      return vector<OperatorDef>{CreateOperatorDef(
                                     "Pow",
                                     "",
                                     std::vector<string>{I(0)},
                                     std::vector<string>{GI(0)},
                                     std::vector<Argument>{pow_arg}),
                                 CreateOperatorDef(
                                     "Mul",
                                     "",
                                     std::vector<string>{GI(0), GO(0)},
                                     std::vector<string>{GI(0)}),
                                 CreateOperatorDef(
                                     "Scale",
                                     "",
                                     std::vector<string>{GI(0)},
                                     std::vector<string>{GI(0)},
                                     std::vector<Argument>{scale_arg})};
      /*
      // Alternative gradient computation
      return vector<OperatorDef>{CreateOperatorDef(
                                     "Div",
                                     "",
                                     std::vector<string>{O(0), I(0)},
                                     std::vector<string>{GI(0)}),
                                 CreateOperatorDef(
                                     "Mul",
                                     "",
                                     std::vector<string>{GI(0), GO(0)},
                                     std::vector<string>{GI(0)}),
                                 CreateOperatorDef(
                                     "Scale",
                                     "",
                                     std::vector<string>{GI(0)},
                                     std::vector<string>{GI(0)},
                                     std::vector<Argument>{scale_arg})};
      */
    } else { // second input is a tensor
      CAFFE_ENFORCE(
          Def().input(0) != Def().output(0) &&
              Def().input(1) != Def().output(0),
          "Gradient computation cannot be carried out if Pow uses in-place "
          "computation: ",
          ProtoDebugString(Def()));
      vector<OperatorDef> grad_ops;
      Argument one_arg;
      one_arg.set_name("value");
      one_arg.set_f(1);
      Argument broadcast, axis, axis_str, order;
      bool bflag = ArgumentHelper::HasArgument(Def(), "broadcast");

      if (bflag) {
        if (ArgumentHelper::HasArgument(Def(), "broadcast")) {
          broadcast = GetArgument(Def(), "broadcast");
        } else {
          broadcast = MakeArgument<int>("broadcast", 0);
        }
        if (ArgumentHelper::HasArgument(Def(), "axis")) {
          axis = GetArgument(Def(), "axis");
        } else {
          axis = MakeArgument<int>("axis", -1);
        }
        if (ArgumentHelper::HasArgument(Def(), "axis_str")) {
          axis_str = GetArgument(Def(), "axis_str");
        } else {
          axis_str = MakeArgument<string>("axis_str", "");
        }
        if (ArgumentHelper::HasArgument(Def(), "order")) {
          order = GetArgument(Def(), "order");
        } else {
          order = MakeArgument<string>("order", "NCHW");
        }
      }

      // function f(w,a) = w^a
      // gradient operator with respect to first input tensor
      // df/dw = a * w^(a-1) (all operations are component-wise)
      grad_ops.push_back(CreateOperatorDef(
          "ConstantFill",
          "",
          std::vector<string>{I(1)},
          std::vector<string>{GI(1)},
          std::vector<Argument>{one_arg}));
      grad_ops.push_back(CreateOperatorDef(
          "Sub",
          "",
          std::vector<string>{I(1), GI(1)},
          std::vector<string>{GI(1)}));
      if (bflag) {
        grad_ops.push_back(CreateOperatorDef(
            "Pow",
            "",
            std::vector<string>{I(0), GI(1)},
            std::vector<string>{GI(0)},
            vector<Argument>{broadcast, axis, axis_str, order}));
      } else {
        grad_ops.push_back(CreateOperatorDef(
            "Pow",
            "",
            std::vector<string>{I(0), GI(1)},
            std::vector<string>{GI(0)}));
      }

      grad_ops.push_back(CreateOperatorDef(
          "Mul",
          "",
          std::vector<string>{GI(0), GO(0)},
          std::vector<string>{GI(0)}));
      if (bflag) {
        grad_ops.push_back(CreateOperatorDef(
            "Mul",
            "",
            std::vector<string>{GI(0), I(1)},
            std::vector<string>{GI(0)},
            vector<Argument>{broadcast, axis, axis_str, order}));
      } else {
        grad_ops.push_back(CreateOperatorDef(
            "Mul",
            "",
            std::vector<string>{GI(0), I(1)},
            std::vector<string>{GI(0)}));
      }
      /*
      // Alternative gradient computation (no broadcast support)
      grad_ops.push_back(CreateOperatorDef(
                           "Div",
                           "",
                           std::vector<string>{O(0), I(0)},
                           std::vector<string>{GI(0)}));
      grad_ops.push_back(CreateOperatorDef(
                           "Mul",
                           "",
                           std::vector<string>{GI(0), GO(0)},
                           std::vector<string>{GI(0)}));
      grad_ops.push_back(CreateOperatorDef(
                           "Mul",
                           "",
                           std::vector<string>{GI(0), I(1)},
                           std::vector<string>{GI(0)}));
      */
      // gradient operator for with respect to second input tensor
      // df/da =  w^a * ln w (all operations are component-wise)
      /*
      // reset GI(1) to zero
      Argument zero_arg;
      zero_arg.set_name("value");
      zero_arg.set_f(0);
      grad_ops.push_back(CreateOperatorDef(
          "ConstantFill",
          "",
          std::vector<string>{I(1)},
          std::vector<string>{GI(1)},
          std::vector<Argument>{zero_arg}));
      */
      grad_ops.push_back(CreateOperatorDef(
          "Log",
          "",
          std::vector<string>{I(0)},
          std::vector<string>{GI(1) + "_autogen_pre_red"}));
      grad_ops.push_back(CreateOperatorDef(
          "Mul",
          "",
          std::vector<string>{GI(1) + "_autogen_pre_red", O(0)},
          std::vector<string>{GI(1) + "_autogen_pre_red"}));
      if (bflag) {
        grad_ops.push_back(CreateOperatorDef(
            "Mul",
            "",
            std::vector<string>{GI(1) + "_autogen_pre_red", GO(0)},
            std::vector<string>{GI(1) + "_autogen_pre_red"}));
        grad_ops.push_back(CreateOperatorDef(
            "SumReduceLike",
            "",
            vector<string>{GI(1) + "_autogen_pre_red", I(1)},
            vector<string>{GI(1)},
            vector<Argument>{axis, axis_str, order}));
      } else {
        grad_ops.push_back(CreateOperatorDef(
            "Mul",
            "",
            std::vector<string>{GI(1) + "_autogen_pre_red", GO(0)},
            std::vector<string>{GI(1)}));
      }

      return grad_ops;
    }
  }

  // Argument `shape` is no longer needed in backprop.
  bool CopyArguments() const override {
    return false;
  }
};

REGISTER_GRADIENT(Pow, GetPowGradient);

} // namespace caffe2