/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file ndarray_function-inl.h
 * \brief The real implementation of NDArray functions.
 */
#ifndef MXNET_NDARRAY_NDARRAY_FUNCTION_INL_H_
#define MXNET_NDARRAY_NDARRAY_FUNCTION_INL_H_

#include <vector>
#include "./ndarray_function.h"
// this file will be included twice by CPU and GPU
// macro to help specialize evaluation function

#ifndef DECL_TERNARY
#define DECL_TERNARY(XPU, OP, FUN)                                                        \
  template <>                                                                             \
  void Eval<XPU, OP>(                                                                     \
      const TBlob& lhs, const TBlob& mhs, const TBlob& rhs, TBlob* ret, RunContext ctx) { \
    FUN<XPU, OP>(lhs, mhs, rhs, ret, ctx);                                                \
  }
#endif

#ifndef DECL_BINARY
#define DECL_BINARY(XPU, OP, FUN)                                                      \
  template <>                                                                          \
  void Eval<XPU, OP>(const TBlob& lhs, const TBlob& rhs, TBlob* ret, RunContext ctx) { \
    FUN<XPU, OP>(lhs, rhs, ret, ctx);                                                  \
  }
#endif

#ifndef DECL_BINARY_LAUNCH
#define DECL_BINARY_LAUNCH(XPU, OP)                                               \
  template <>                                                                     \
  void BinaryOpKernelImpl<OP, XPU>(                                               \
      mshadow::Stream<XPU> * s, const TBlob& lhs, const TBlob& rhs, TBlob* out) { \
    BinaryOpKernelLaunch<OP>(s, lhs, rhs, out);                                   \
  }
#endif

#ifndef DECL_SCALAR
#define DECL_SCALAR(XPU, OP, FUN, REVERSE)                                                       \
  template <>                                                                                    \
  void Eval<XPU, OP, REVERSE>(const TBlob& lhs, const real_t& rhs, TBlob* ret, RunContext ctx) { \
    FUN<XPU, OP, REVERSE>(lhs, rhs, ret, ctx);                                                   \
  }
#endif

#if defined(__CUDACC__)
#define DEVICE gpu
#else
#define DEVICE cpu
#endif

namespace mxnet {
namespace ndarray {

// true implementation
template <typename xpu, typename OP>
void EvalBinary_(const TBlob& lhs, const TBlob& rhs, TBlob* ret, RunContext ctx) {
  using namespace mshadow::expr;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  CHECK_EQ(ret->type_flag_, lhs.type_flag_) << "Only support input/output with the same data type";
  CHECK_EQ(ret->type_flag_, rhs.type_flag_) << "Only support input/output with the same data type";
  MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
    ret->FlatTo2D<xpu, DType>(s) =
        F<typename OP::mshadow_op>(lhs.FlatTo2D<xpu, DType>(s), rhs.FlatTo2D<xpu, DType>(s));
  });
}

template <typename xpu, typename OP>
void EvalOneHot_(const TBlob& index, const TBlob& rhs, TBlob* ret, RunContext ctx) {
  LOG(INFO) << "The operator onehot_encode is deprecated; use one_hot instead.";
  using namespace mshadow::expr;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  // TODO(eric): support mixed type encoding, i.e. int index and float rhs.
  CHECK_EQ(ret->type_flag_, mshadow::default_type_flag)
      << "one_hot_encode only support float32 as input/output";
  CHECK_EQ(rhs.type_flag_, mshadow::default_type_flag)
      << "one_hot_encode only support float32 as input/output";
  CHECK_EQ(index.type_flag_, mshadow::default_type_flag)
      << "one_hot_encode only support float32 as input/output";
  ret->get<xpu, 2, real_t>(s) = one_hot_encode(index.get<xpu, 1, real_t>(s), rhs.shape_[1]);
}

template <typename xpu, typename OP>
void EvalMatChooseRowElem_(const TBlob& lhs, const TBlob& rhs, TBlob* ret, RunContext ctx) {
  using namespace mshadow::expr;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  // TODO(eric): support mixed type choose, i.e. int index and float rhs.
  CHECK_EQ(ret->type_flag_, mshadow::default_type_flag)
      << "mat_choose_row_element only support float32 as input/output";
  CHECK_EQ(rhs.type_flag_, mshadow::default_type_flag)
      << "mat_choose_row_element only support float32 as input/output";
  CHECK_EQ(lhs.type_flag_, mshadow::default_type_flag)
      << "mat_choose_row_element only support float32 as input/output";
  ret->get<xpu, 1, real_t>(s) =
      mat_choose_row_element(lhs.get<xpu, 2, real_t>(s), rhs.get<xpu, 1, real_t>(s));
}

template <typename xpu, typename OP>
void EvalMatFillRowElem_(const TBlob& lhs,
                         const TBlob& mhs,
                         const TBlob& rhs,
                         TBlob* ret,
                         RunContext ctx) {
  using namespace mshadow::expr;
  mshadow::Stream<xpu>* s     = ctx.get_stream<xpu>();
  ret->get<xpu, 2, real_t>(s) = mat_fill_row_element(
      lhs.get<xpu, 2, real_t>(s), mhs.get<xpu, 1, real_t>(s), rhs.get<xpu, 1, real_t>(s));
}

template <typename xpu, typename OP, bool reverse>
void EvalScalar_(const TBlob& lhs, const real_t& rhs, TBlob* ret, RunContext ctx) {
  using namespace mshadow::expr;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  CHECK_EQ(ret->type_flag_, lhs.type_flag_) << "Only support input/output with the same data type";
  if (reverse) {
    MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
      ret->FlatTo2D<xpu, DType>(s) =
          F<typename OP::mshadow_op>(scalar(DType(rhs)), lhs.FlatTo2D<xpu, DType>(s));
    });
  } else {
    MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
      ret->FlatTo2D<xpu, DType>(s) =
          F<typename OP::mshadow_op>(lhs.FlatTo2D<xpu, DType>(s), scalar(DType(rhs)));
    });
  }
}

template <>
void EvalClip<DEVICE>(const TBlob& src,
                      const real_t& a_min,
                      const real_t& a_max,
                      TBlob* ret,
                      RunContext ctx) {
  typedef DEVICE xpu;
  using namespace mshadow::expr;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  CHECK_EQ(ret->type_flag_, src.type_flag_) << "Only support input/output with the same data type";
  MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
    ret->FlatTo2D<xpu, DType>(s) = F<ClipMax::mshadow_op>(
        F<ClipMin::mshadow_op>(src.FlatTo2D<xpu, DType>(s), scalar(DType(a_min))),
        scalar(DType(a_max)));
  });
}

template <>
void EvalRandom<DEVICE, UniformDistribution>(const real_t& a,
                                             const real_t& b,
                                             const Resource& resource,
                                             TBlob* ret,
                                             RunContext ctx) {
  typedef DEVICE xpu;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SampleUniform(&tmp, float(a), float(b));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SampleUniform(&tmp, double(a), double(b));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

template <>
void EvalRandom<DEVICE, GaussianDistribution>(const real_t& mu,
                                              const real_t& sigma,
                                              const Resource& resource,
                                              TBlob* ret,
                                              RunContext ctx) {
  typedef DEVICE xpu;
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SampleGaussian(&tmp, float(mu), float(sigma));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SampleGaussian(&tmp, double(mu), double(sigma));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

#if defined(__CUDACC__)

template <>
void EvalRandom<gpu, GammaDistribution>(const real_t& alpha,
                                        const real_t& beta,
                                        const Resource& resource,
                                        TBlob* ret,
                                        RunContext ctx) {
  EvalRandom<cpu, GammaDistribution>(alpha, beta, resource, ret, ctx);
}

template <>
void EvalRandom<gpu, ExponentialDistribution>(
    const real_t& lambda,
    const real_t& dummy,  // this is to satisfy the SampleOp lambda signature
    const Resource& resource,
    TBlob* ret,
    RunContext ctx) {
  EvalRandom<cpu, ExponentialDistribution>(lambda, dummy, resource, ret, ctx);
}

template <>
void EvalRandom<gpu, PoissonDistribution>(
    const real_t& lambda,
    const real_t& dummy,  // this is to satisfy the SampleOp lambda signature
    const Resource& resource,
    TBlob* ret,
    RunContext ctx) {
  EvalRandom<cpu, PoissonDistribution>(lambda, dummy, resource, ret, ctx);
}

template <>
void EvalRandom<gpu, NegBinomialDistribution>(const real_t& k,
                                              const real_t& p,
                                              const Resource& resource,
                                              TBlob* ret,
                                              RunContext ctx) {
  EvalRandom<cpu, NegBinomialDistribution>(k, p, resource, ret, ctx);
}

template <>
void EvalRandom<gpu, GenNegBinomialDistribution>(const real_t& mu,
                                                 const real_t& alpha,
                                                 const Resource& resource,
                                                 TBlob* ret,
                                                 RunContext ctx) {
  EvalRandom<cpu, GenNegBinomialDistribution>(mu, alpha, resource, ret, ctx);
}

#else

template <>
void EvalRandom<cpu, GammaDistribution>(const real_t& alpha,
                                        const real_t& beta,
                                        const Resource& resource,
                                        TBlob* ret,
                                        RunContext ctx) {
  typedef cpu xpu;  // No support for gpu for this distribution.
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SampleGamma(&tmp, float(alpha), float(beta));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SampleGamma(&tmp, double(alpha), double(beta));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

template <>
void EvalRandom<cpu, ExponentialDistribution>(
    const real_t& lambda,
    const real_t& dummy,  // this is to satisfy the SampleOp lambda signature
    const Resource& resource,
    TBlob* ret,
    RunContext ctx) {
  typedef cpu xpu;  // No support for gpu for this distribution.
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SampleExponential(&tmp, float(lambda));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SampleExponential(&tmp, double(lambda));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

template <>
void EvalRandom<cpu, PoissonDistribution>(
    const real_t& lambda,
    const real_t& dummy,  // this is to satisfy the SampleOp lambda signature
    const Resource& resource,
    TBlob* ret,
    RunContext ctx) {
  typedef cpu xpu;  // No support for gpu for this distribution.
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SamplePoisson(&tmp, float(lambda));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SamplePoisson(&tmp, double(lambda));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

template <>
void EvalRandom<cpu, NegBinomialDistribution>(const real_t& k,
                                              const real_t& p,
                                              const Resource& resource,
                                              TBlob* ret,
                                              RunContext ctx) {
  typedef cpu xpu;  // No support for gpu for this distribution.
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SampleNegativeBinomial(&tmp, float(k), float(p));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SampleNegativeBinomial(&tmp, double(k), double(p));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

template <>
void EvalRandom<cpu, GenNegBinomialDistribution>(const real_t& mu,
                                                 const real_t& alpha,
                                                 const Resource& resource,
                                                 TBlob* ret,
                                                 RunContext ctx) {
  typedef cpu xpu;  // No support for gpu for this distribution.
  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
  switch (ret->type_flag_) {
    case mshadow::kFloat32: {
      mshadow::Random<xpu, float>* prnd  = resource.get_random<xpu, float>(s);
      mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s);
      prnd->SampleGeneralizedNegativeBinomial(&tmp, float(mu), float(alpha));  // NOLINT(*)
      break;
    }
    case mshadow::kFloat64: {
      mshadow::Random<xpu, double>* prnd  = resource.get_random<xpu, double>(s);
      mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s);
      prnd->SampleGeneralizedNegativeBinomial(&tmp, double(mu), double(alpha));  // NOLINT(*)
      break;
    }
    default:
      LOG(FATAL) << "Random only support float32 and float64";
  }
}

#endif  // #ifndef __CUDACC__

template <>
void Eval<DEVICE>(const real_t& rhs, TBlob* ret, RunContext ctx) {
  mshadow::Stream<DEVICE>* s = ctx.get_stream<DEVICE>();
  MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(
      ret->type_flag_, DType, { ret->FlatTo2D<DEVICE, DType>(s) = DType(rhs); });
}

template <>
void ElementwiseSum<DEVICE>(const std::vector<TBlob> source, TBlob* dst, RunContext ctx) {
  typedef DEVICE xpu;
  using namespace mshadow;
  using namespace mshadow::expr;
  Stream<xpu>* s = ctx.get_stream<xpu>();
  for (size_t i = 1; i < source.size(); ++i) {
    CHECK_EQ(source[i].type_flag_, dst->type_flag_)
        << "Only support input/output with the same data type";
  }
  MSHADOW_TYPE_SWITCH(dst->type_flag_, DType, {
    Tensor<xpu, 2, DType> out = dst->FlatTo2D<xpu, DType>(s);

    switch (source.size()) {
      case 2: {
        Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s);
        Tensor<xpu, 2, DType> in_1 = source[1].FlatTo2D<xpu, DType>(s);
        out                        = in_0 + in_1;
        break;
      }
      case 3: {
        Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s);
        Tensor<xpu, 2, DType> in_1 = source[1].FlatTo2D<xpu, DType>(s);
        Tensor<xpu, 2, DType> in_2 = source[2].FlatTo2D<xpu, DType>(s);
        out                        = in_0 + in_1 + in_2;
        break;
      }
      case 4: {
        Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s);
        Tensor<xpu, 2, DType> in_1 = source[1].FlatTo2D<xpu, DType>(s);
        Tensor<xpu, 2, DType> in_2 = source[2].FlatTo2D<xpu, DType>(s);
        Tensor<xpu, 2, DType> in_3 = source[3].FlatTo2D<xpu, DType>(s);
        out                        = in_0 + in_1 + in_2 + in_3;
        break;
      }
      default: {
        Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s);
        out                        = F<op::mshadow_op::identity>(in_0);
        for (size_t i = 1; i < source.size(); ++i) {
          out += source[i].FlatTo2D<xpu, DType>(s);
        }
        break;
      }
    }
  });
}

template <>
void EvalBroadcast<DEVICE>(TBlob const& src, TBlob* ret, int size, RunContext ctx) {
  typedef DEVICE xpu;
  mshadow::Stream<xpu>* s     = ctx.get_stream<xpu>();
  mshadow::Tensor<xpu, 3> out = ret->get<xpu, 3, real_t>(s);
  mshadow::Tensor<xpu, 2> in  = src.get<xpu, 2, real_t>(s);
  out                         = mshadow::expr::broadcast_with_axis(in, 0, size);
}

template <typename OP, typename xpu>
void BinaryOpKernelLaunch(mshadow::Stream<xpu>* s, const TBlob& lhs, const TBlob& rhs, TBlob* out) {
  using namespace op::mxnet_op;
  using namespace mshadow;
  MSHADOW_TYPE_SWITCH(out->type_flag_, DType, {
    Kernel<op_with_req<OP, kWriteInplace>, xpu>::Launch(
        s, lhs.Size(), out->dptr<DType>(), lhs.dptr<DType>(), rhs.dptr<DType>());
  });
}
// declarations
DECL_BINARY(DEVICE, MatChooseRowElem, EvalMatChooseRowElem_)
DECL_TERNARY(DEVICE, MatFillRowElem, EvalMatFillRowElem_)
DECL_BINARY(DEVICE, OneHotEncode, EvalOneHot_)
DECL_SCALAR(DEVICE, Plus, EvalScalar_, true)
DECL_SCALAR(DEVICE, Minus, EvalScalar_, true)
DECL_SCALAR(DEVICE, Mul, EvalScalar_, true)
DECL_SCALAR(DEVICE, Div, EvalScalar_, true)
DECL_BINARY_LAUNCH(DEVICE, Plus)
DECL_BINARY_LAUNCH(DEVICE, Minus)
DECL_BINARY_LAUNCH(DEVICE, Mul)
DECL_BINARY_LAUNCH(DEVICE, Div)

// for reverse seq
DECL_SCALAR(DEVICE, Plus, EvalScalar_, false)
DECL_SCALAR(DEVICE, Minus, EvalScalar_, false)
DECL_SCALAR(DEVICE, Mul, EvalScalar_, false)
DECL_SCALAR(DEVICE, Div, EvalScalar_, false)
}  // namespace ndarray
}  // namespace mxnet

#endif  // MXNET_NDARRAY_NDARRAY_FUNCTION_INL_H_