/*
 * 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.cc
 * \brief CPU Implementation of ndarray function.
 */

// this will be invoked by gcc and compile CPU version
#include "./ndarray_function.h"
#include "./ndarray_function-inl.h"
#include "../common/utils.h"
#include "../operator/mxnet_op.h"
#include "../operator/tensor/elemwise_binary_op-inl.h"
#include "../operator/tensor/elemwise_sum.h"

namespace mxnet {
namespace ndarray {
template <>
void Copy<cpu, cpu>(const TBlob& from,
                    TBlob* to,
                    Context from_ctx,
                    Context to_ctx,
                    RunContext ctx) {
  MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(to->type_flag_, DType, {
    if (to->type_flag_ == from.type_flag_) {
      if (!features::is_enabled(features::INT64_TENSOR_SIZE)) {
        CHECK_LT(from.Size(), (int64_t{1} << 31) - 1)
            << "Size of tensor you are trying to allocate is larger than "
               "2^31 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
      }
      const index_t size = static_cast<index_t>(from.Size());
      CHECK_EQ(size, to->Size()) << "copying size mismatch, from: " << size * sizeof(DType)
                                 << " bytes, to: " << to->Size() * sizeof(DType) << " bytes.";
      common::ParallelCopy(to->dptr<DType>(), from.dptr<DType>(), size);
    } else {
      MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(from.type_flag_, SrcDType, {
        to->FlatTo1D<cpu, DType>() = mshadow::expr::tcast<DType>(from.FlatTo1D<cpu, SrcDType>());
      })
    }
  })
}

template <typename DType, typename IType>
void ElementwiseSumRspImpl(mshadow::Stream<cpu>* s,
                           const std::vector<NDArray>& nds,
                           const std::vector<IType>& uniq_row_idx,
                           NDArray* out,
                           const int nthreads = 4) {
#pragma omp parallel num_threads(nthreads)
  {
    const size_t nnr             = uniq_row_idx.size();
    const int num_threads        = omp_get_num_threads();
    size_t row_block_len         = (nnr + num_threads - 1) / num_threads;
    const size_t row_block_start = omp_get_thread_num() * row_block_len;
    if (row_block_start < nnr) {
      const size_t row_block_end = std::min(row_block_start + row_block_len, nnr);

      const size_t row_length = out->data().shape_.ProdShape(1, out->data().shape_.ndim());
      auto out_values         = out->data().get_with_shape<cpu, 2, DType>(
          mshadow::Shape2(out->storage_shape()[0], row_length), s);
      auto out_indices = out->aux_data(rowsparse::kIdx).FlatTo1D<cpu, IType>();
      for (size_t i = row_block_start; i < row_block_end; ++i) {
        out_indices[i] = uniq_row_idx[i];
      }
      for (const auto& nd : nds) {
        if (nd.storage_initialized()) {
          const auto nd_indices = nd.aux_data(rowsparse::kIdx).FlatTo1D<cpu, IType>();
          const auto nd_values  = nd.data().get_with_shape<cpu, 2, DType>(
              mshadow::Shape2(nd.storage_shape()[0], row_length), s);
          const auto nd_num_rows        = nd.aux_shape(rowsparse::kIdx).Size();
          const IType* nd_indices_start = &nd_indices[0];
          const IType* nd_indices_end   = nd_indices_start + nd_num_rows;
          const IType* row_idx_ptr =
              std::lower_bound(nd_indices_start, nd_indices_end, out_indices[row_block_start]);
          // skip this nd if all of its row indices are smaller than out_indices[row_block_start]
          // or current row block is not covered by [*row_idx_ptr, nd_indices_end).
          if (nd_indices_end == row_idx_ptr || *row_idx_ptr > out_indices[row_block_end - 1]) {
            continue;
          }
          for (size_t irow = row_block_start;
               irow < row_block_end && row_idx_ptr != nd_indices_end;) {
            if (out_indices[irow] == *row_idx_ptr) {
              auto out_value_cur_row = out_values[irow];
              const auto offset      = row_idx_ptr - nd_indices_start;
              auto nd_value_cur_row  = nd_values[offset];
              for (index_t j = 0; j < nd_value_cur_row.shape_[0]; ++j) {
                out_value_cur_row[j] += nd_value_cur_row[j];
              }
              ++irow;
              ++row_idx_ptr;
            } else if (out_indices[irow] < *row_idx_ptr) {
              ++irow;
            } else {
              ++row_idx_ptr;
            }
          }
        }
      }
    }
  }
}

/*!
 * \brief Given a vector of ndarrays, generate a index vector containing
 * all the unique row indices of the ndarrays.
 */
template <typename IType>
void GetUniqueRspRowIdx(const std::vector<NDArray>& nds, std::vector<IType>* uniq_row_idx) {
  using namespace rowsparse;
  size_t total_num_rows = 0;
  for (const auto& nd : nds) {
    CHECK_EQ(nd.storage_type(), kRowSparseStorage);
    if (nd.storage_initialized()) {
      total_num_rows += nd.aux_shape(kIdx).Size();
    }
  }

  uniq_row_idx->resize(total_num_rows);
  int nthreads = omp_get_max_threads();
  int offset   = 0;
  for (const auto& nd : nds) {
    if (nd.storage_initialized()) {
      const IType* nd_row_idx = nd.aux_data(kIdx).dptr<IType>();
      const int num_rows      = nd.aux_shape(kIdx).Size();
#pragma omp parallel for num_threads(nthreads)
      for (int i = 0; i < num_rows; ++i) {
        (*uniq_row_idx)[offset + i] = nd_row_idx[i];
      }
      offset += num_rows;
    }
  }

  common::ParallelSort(uniq_row_idx->begin(), uniq_row_idx->end(), nthreads);
  auto it = std::unique(uniq_row_idx->begin(), uniq_row_idx->end());
  uniq_row_idx->resize(it - uniq_row_idx->begin());
}

void ElementwiseSumRsp(mshadow::Stream<cpu>* s,
                       const Resource& rsc,
                       const std::vector<NDArray>& nds,
                       NDArray* out) {
  if (nds.empty())
    return;
  using namespace rowsparse;
  CHECK_EQ(out->storage_type(), kRowSparseStorage)
      << "Expected row sparse storage type (" << out->storage_type() << " given)";

  MSHADOW_TYPE_SWITCH(out->dtype(), DType, {
    MSHADOW_IDX_TYPE_SWITCH(out->aux_type(kIdx), IType, {
      // TODO(Jun): Use resource rsc for temporary vector instead of
      //            allocating it directly in GetUniqueRspRowIdx
      std::vector<IType> uniq_row_idx;
      GetUniqueRspRowIdx(nds, &uniq_row_idx);
      out->CheckAndAlloc({mshadow::Shape1(uniq_row_idx.size())});
      out->data().FlatTo2D<cpu, DType>() = static_cast<DType>(0);
      ElementwiseSumRspImpl<DType, IType>(s, nds, uniq_row_idx, out, omp_get_max_threads());
    });
  });
}

void ElementwiseSumDnsCsrDnsImpl(mshadow::Stream<cpu>* s,
                                 const Resource& rsc,
                                 const std::vector<NDArray>& nds,
                                 NDArray* out) {
  using namespace mxnet::op;
  using namespace mxnet::op::mxnet_op;
  const TBlob& out_data = out->data();
  MSHADOW_TYPE_SWITCH(out->dtype(), DType, {  // data type
    Kernel<Sum, cpu>::Launch(s,
                             out_data.Size(),
                             out_data.dptr<DType>(),
                             kWriteTo,
                             nds[0].data().dptr<DType>(),
                             nds[2].data().dptr<DType>());
    const TBlob& csr_data      = nds[1].data();
    const TBlob& csr_indices   = nds[1].aux_data(csr::kIdx);
    const TBlob& csr_indptr    = nds[1].aux_data(csr::kIndPtr);
    const nnvm::dim_t num_rows = nds[1].shape()[0];
    const nnvm::dim_t num_cols = nds[1].shape()[1];
    MSHADOW_IDX_TYPE_SWITCH(csr_indices.type_flag_, IType, {   // indices type
      MSHADOW_IDX_TYPE_SWITCH(csr_indptr.type_flag_, CType, {  // indptr type
        if (nds[1].storage_initialized()) {
          Kernel<ElemwiseDnsCsrDnsKernel<kWriteTo, mshadow_op::plus>, cpu>::Launch(
              s,
              num_rows,
              out_data.dptr<DType>(),
              out_data.dptr<DType>(),
              csr_data.dptr<DType>(),
              csr_indices.dptr<IType>(),
              csr_indptr.dptr<CType>(),
              num_rows,
              num_cols);
        }
      });
    });
  });
}

void ElementwiseSumContainsDnsImpl(mshadow::Stream<cpu>* s,
                                   const Resource& rsc,
                                   const std::vector<NDArray>& nds,
                                   NDArray* out) {
  using namespace mxnet::op;
  using namespace mxnet::op::mxnet_op;
  const TBlob& out_data = out->data();
  MSHADOW_TYPE_SWITCH(out->dtype(), DType, {  // data type
    // Do not set_zero when output mem inplace with input[0] mem
    // Now for add_n OP, output mem can be in-placed with the first input
    if (nds[0].data().dptr<DType>() != out_data.dptr<DType>()) {
      Kernel<set_zero, cpu>::Launch(s, out_data.Size(), out_data.dptr<DType>());
    }
    for (size_t i = 0; i < nds.size(); ++i) {
      const NDArray& nd    = nds[i];
      const TBlob& nd_data = nd.data();

      if (i == 0) {
        if (nd.storage_type() == kDefaultStorage) {
          Kernel<op_with_req<mshadow_op::identity, kWriteTo>, cpu>::Launch(
              s, out_data.Size(), out_data.dptr<DType>(), nd_data.dptr<DType>());
          continue;
        } else {
          Kernel<set_zero, cpu>::Launch(s, out_data.Size(), out_data.dptr<DType>());
        }
      }

      switch (nd.storage_type()) {
        case kDefaultStorage: {
          Kernel<op_with_req<mshadow_op::plus, kWriteTo>, cpu>::Launch(s,
                                                                       out_data.Size(),
                                                                       out_data.dptr<DType>(),
                                                                       out_data.dptr<DType>(),
                                                                       nd_data.dptr<DType>());
          break;
        }
        case kCSRStorage: {
          const TBlob& nd_indices    = nd.aux_data(csr::kIdx);
          const TBlob& nd_indptr     = nd.aux_data(csr::kIndPtr);
          const nnvm::dim_t num_rows = nd.shape()[0];
          const nnvm::dim_t num_cols = nd.shape()[1];
          MSHADOW_IDX_TYPE_SWITCH(nd_indices.type_flag_, IType, {   // indices type
            MSHADOW_IDX_TYPE_SWITCH(nd_indptr.type_flag_, CType, {  // indptr type
              if (nd.storage_initialized()) {
                Kernel<ElemwiseDnsCsrDnsKernel<kWriteTo, mshadow_op::plus>, cpu>::Launch(
                    s,
                    num_rows,
                    out_data.dptr<DType>(),
                    out_data.dptr<DType>(),
                    nd_data.dptr<DType>(),
                    nd_indices.dptr<IType>(),
                    nd_indptr.dptr<CType>(),
                    num_rows,
                    num_cols);
              }
            });
          });
          break;
        }
        case kRowSparseStorage: {
          const TBlob& nd_indices    = nd.aux_data(rowsparse::kIdx);
          const nnvm::dim_t num_rows = nd.shape()[0];
          const nnvm::dim_t num_cols = nd.shape()[1];
          MSHADOW_IDX_TYPE_SWITCH(nd_indices.type_flag_, IType, {  // indices type
            if (nd.storage_initialized()) {
              const nnvm::dim_t nz_rows = nd_indices.Size();
              Kernel<ElemwiseDnsRspDnsKernel<kWriteTo, mshadow_op::plus>, cpu>::Launch(
                  s,
                  nz_rows * num_cols,
                  out_data.dptr<DType>(),
                  out_data.dptr<DType>(),
                  nd_data.dptr<DType>(),
                  nd_indices.dptr<IType>(),
                  num_rows,
                  nz_rows,
                  num_cols);
            }
          });
          break;
        }
        default:
          LOG(FATAL) << "unknown storage type " << nd.storage_type() << "encountered...";
      }
    }
  });
}

/*!
 * \brief Parallel cpu impl of elemwise sum for sparse tensors.
 * Currently only support row sparse sum.
 */
template <>
void ElementwiseSum<cpu>(mshadow::Stream<cpu>* s,
                         const Resource& rsc,
                         const std::vector<NDArray>& nds,
                         NDArray* out) {
  if (nds.empty())
    return;
  if (common::ContainsOnlyStorage(nds, kRowSparseStorage)) {
    ElementwiseSumRsp(s, rsc, nds, out);
  } else if (nds.size() == 3U && nds[0].storage_type() == kDefaultStorage &&
             nds[1].storage_type() == kCSRStorage && nds[2].storage_type() == kDefaultStorage &&
             out->storage_type() == kDefaultStorage) {
    ElementwiseSumDnsCsrDnsImpl(s, rsc, nds, out);
  } else if (nds.size() > 4U && common::ContainsStorageType(nds, kDefaultStorage) &&
             out->storage_type() == kDefaultStorage) {
    ElementwiseSumContainsDnsImpl(s, rsc, nds, out);
  } else {
    LOG(FATAL) << "ElementwiseSum<cpu> has not been implemented for storage_type = << "
               << nds[0].storage_type();
  }
}

template <>
void Eval<cpu>(mshadow::Stream<cpu>* s, const real_t val, const NDArray& dst) {
  NDArray temp                   = dst;
  const NDArrayStorageType stype = temp.storage_type();
  if (stype == kRowSparseStorage) {
    SetValueRspImpl(s, val, &temp);
  } else {
    LOG(FATAL) << "Not implemented for storage type" << stype;
  }
}

}  // namespace ndarray
}  // namespace mxnet