/*
 * 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 np_unique_op.cc
 */

#include "./np_unique_op.h"

namespace mxnet {
namespace op {

inline bool NumpyUniqueType(const nnvm::NodeAttrs& attrs,
                            std::vector<int> *in_attrs,
                            std::vector<int> *out_attrs) {
  CHECK_EQ(in_attrs->size(), 1U);
  TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
  TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
  for (size_t i = 1; i < out_attrs->size(); ++i) {
    TYPE_ASSIGN_CHECK(*out_attrs, i, mshadow::kInt64);
  }
  return out_attrs->at(0) != -1;
}

inline bool NumpyUniqueStorageType(const nnvm::NodeAttrs& attrs,
                                   const int dev_mask,
                                   DispatchMode* dispatch_mode,
                                   std::vector<int> *in_attrs,
                                   std::vector<int> *out_attrs) {
  CHECK_EQ(in_attrs->size(), 1U);
  // CHECK_EQ(out_attrs->size(), 1U);
  for (int &attr : *in_attrs) {
    CHECK_EQ(attr, kDefaultStorage) << "Only default storage is supported";
  }
  for (int &attr : *out_attrs) {
    attr = kDefaultStorage;
  }
  *dispatch_mode = DispatchMode::kFComputeEx;
  return true;
}

struct UniqueComputeAuxCPUKernel {
  // assume that idx have been flattened to a 1-D tensor (N,)
  // assume that out_data and in_data have been flattened to 2-D tensors, (N, M) and (K, M)
  // M is the number of columns of in_data and out_data
  // i is the index of out_data
  template<typename DType>
  MSHADOW_XINLINE static void Map(dim_t i, DType* out_data, const DType* in_data,
                                  const dim_t* idx, const dim_t M) {
    dim_t j = idx[i];
    std::memcpy(out_data + i * M, in_data + j * M, M * sizeof(DType));
  }
};

struct UniqueComputeMaskCPUKernel {
  template<typename DType>
  MSHADOW_XINLINE static void Map(dim_t i,
                                  dim_t* out_data,
                                  const DType* in_data,
                                  const dim_t numel) {
    if (i == 0) {
      out_data[i] = 1;
    } else {
      out_data[i] = (std::memcmp(in_data + i * numel,
                     in_data + (i - 1) * numel, numel * sizeof(DType)) == 0) ? 0 : 1;
    }
  }
};

void NumpyUniqueCPUNoneAxisImpl(const NumpyUniqueParam& param,
                                const OpContext &ctx,
                                const std::vector<NDArray> &inputs,
                                const std::vector<OpReqType> &req,
                                const std::vector<NDArray> &outputs) {
  MSHADOW_TYPE_SWITCH(outputs[0].dtype(), DType, {
    mshadow::Stream<cpu> *stream = ctx.get_stream<cpu>();

    DType* input_data = inputs[0].data().dptr<DType>();
    dim_t input_size = inputs[0].shape().Size();
    if (param.return_index || param.return_inverse || param.return_counts) {
      // argsort, result in perm
      std::vector<dim_t> perm(input_size);
      std::iota(perm.begin(), perm.end(), 0);
      std::stable_sort(perm.begin(), perm.end(),
        [&input_data] (dim_t i1, dim_t i2) {return input_data[i1] < input_data[i2];});
      // sorted data in aux
      std::vector<DType> aux(input_size);
      mxnet_op::Kernel<UniqueComputeAuxCPUKernel, cpu>::Launch(
        stream, input_size, aux.data(), input_data, perm.data(), 1);
      // calculate unique mask
      std::vector<dim_t> mask(input_size);
      mxnet_op::Kernel<UniqueComputeMaskCPUKernel, cpu>::Launch(
        stream, input_size, mask.data(), aux.data(), 1);
      // Calculate prefix sum
      std::vector<int32_t> prefix_sum(input_size, 0);
      int32_t valid_num = 0;
      for (dim_t i = 0; i < input_size; i++) {
        prefix_sum[i] = (i == 0) ? 0 : prefix_sum[i - 1];
        prefix_sum[i] += (mask[i]) ? 1 : 0;
      }
      valid_num = prefix_sum[input_size - 1];
      // set the output shape forcefully
      mxnet::TShape s(1, valid_num);
      const_cast<NDArray &>(outputs[0]).Init(s);
      // launch kernal to obtain unique array, reuse boolean_mask kernel
      mxnet_op::Kernel<BooleanMaskForwardCPUKernel, cpu>::Launch(
        stream, input_size, outputs[0].data().dptr<DType>(), aux.data(),
        prefix_sum.data(), 1);
      // handle other optional outputs
      int output_flag = 0;
      if (param.return_index) {
        output_flag += 1;
        const_cast<NDArray &>(outputs[output_flag]).Init(s);
        dim_t* unique_indices = outputs[output_flag].data().dptr<dim_t>();
        // reuse boolean_mask kernel
        mxnet_op::Kernel<BooleanMaskForwardCPUKernel, cpu>::Launch(
          stream, input_size, unique_indices, perm.data(),
          prefix_sum.data(), 1);
      }
      if (param.return_inverse) {
        output_flag += 1;
        const_cast<NDArray &>(outputs[output_flag]).Init(mxnet::TShape(1, input_size));
        dim_t* unique_inverse = outputs[output_flag].data().dptr<dim_t>();
        mxnet_op::Kernel<UniqueReturnInverseKernel, cpu>::Launch(
          stream, input_size, unique_inverse, prefix_sum.data(), perm.data());
      }
      if (param.return_counts) {
        output_flag += 1;
        std::vector<dim_t> idx(valid_num + 1);
        auto iter = idx.begin();
        for (dim_t i = 0; i < input_size; ++i) {
          if (mask[i]) {
            *iter = i;
            ++iter;
          }
        }
        *iter = input_size;
        const_cast<NDArray &>(outputs[output_flag]).Init(s);
        dim_t* unique_counts = outputs[output_flag].data().dptr<dim_t>();
        mxnet_op::Kernel<UniqueReturnCountsKernel, cpu>::Launch(
          stream, valid_num, unique_counts, idx.data());
      }
    } else {
      std::set<DType> set(input_data, input_data + input_size);
      mxnet::TShape s(1, set.size());
      const_cast<NDArray &>(outputs[0]).Init(s);
      std::copy(set.begin(), set.end(), outputs[0].data().dptr<DType>());
    }
  });
}

void NumpyUniqueCPUImpl(const NumpyUniqueParam& param,
                        const OpContext &ctx,
                        const std::vector<NDArray> &inputs,
                        const std::vector<OpReqType> &req,
                        const std::vector<NDArray> &outputs) {
  CHECK(param.axis.value() >= -1 * inputs[0].shape().ndim() &&
      param.axis.value() < inputs[0].shape().ndim())
      << "Axis should be in the range of [-r, r-1] where r is the rank of input tensor";
  MSHADOW_TYPE_SWITCH(outputs[0].dtype(), DType, {
    using namespace mshadow;
    using namespace mshadow::expr;
    Stream<cpu> *stream = ctx.get_stream<cpu>();
    const index_t actual_axis =
        param.axis.value() + ((param.axis.value() < 0) ? inputs[0].shape().ndim() : 0);
    // reshape tensor to [origin_shape[axis], -1]
    const mxnet::TShape origin_shape = inputs[0].shape();
    Tensor<cpu, 3, DType> input_tensor_3d =
        inputs[0].data().FlatTo3D<cpu, DType>(actual_axis, stream);
    Tensor<cpu, 1, DType> workspace = ctx.requested[0].get_space_typed<cpu, 1, DType>(
        Shape1(input_tensor_3d.shape_.Size() * 2), stream);
    Tensor<cpu, 3, DType> input_tensor(workspace.dptr_, Shape3(
        input_tensor_3d.shape_[1], input_tensor_3d.shape_[0], input_tensor_3d.shape_[2]), stream);
    input_tensor = swapaxis<1, 0>(input_tensor_3d);
    const Shape<3> temp_shape = input_tensor.shape_;
    DType* input_data = input_tensor.dptr_;
    dim_t numel = temp_shape[1] * temp_shape[2];
    // argsort, result in perm
    std::vector<dim_t> perm(temp_shape[0]);
    std::iota(perm.begin(), perm.end(), 0);
    std::stable_sort(perm.begin(), perm.end(),
      [&](dim_t a, dim_t b) -> bool {
        for (dim_t i = 0; i < numel; ++i) {
          DType lhs = input_data[i + a * numel];
          DType rhs = input_data[i + b * numel];
          if (lhs < rhs) {
            return true;
          } else if (lhs > rhs) {
            return false;
          }
        }
        return false;
      });
    // sorted data in aux
    Tensor<cpu, 2, DType> aux(workspace.dptr_ + input_tensor_3d.shape_.Size(),
        Shape2(temp_shape[0], temp_shape[1] * temp_shape[2]), stream);
    mxnet_op::Kernel<UniqueComputeAuxCPUKernel, cpu>::Launch(
      stream, temp_shape[0], aux.dptr_, input_data, perm.data(), numel);
    // calculate unique mask
    std::vector<dim_t> mask(temp_shape[0]);
    mxnet_op::Kernel<UniqueComputeMaskCPUKernel, cpu>::Launch(
      stream, temp_shape[0], mask.data(), aux.dptr_, numel);
    // calculate prefix sum
    std::vector<int32_t> prefix_sum(temp_shape[0], 0);
    int32_t valid_num = 0;
    for (dim_t i = 0; i < temp_shape[0]; i++) {
      prefix_sum[i] = (i == 0) ? 0 : prefix_sum[i - 1];
      prefix_sum[i] += (mask[i]) ? 1 : 0;
    }
    valid_num = prefix_sum[temp_shape[0] - 1];
    // store the temp output data, reuse the space of 'input_tensor'
    Tensor<cpu, 3, DType> temp_tensor(workspace.dptr_,
        Shape3(valid_num, temp_shape[1], temp_shape[2]), stream);
    // launch kernal to obtain unique array, reuse boolean_mask kernel
    mxnet_op::Kernel<BooleanMaskForwardCPUKernel, cpu>::Launch(
      stream, temp_shape[0], temp_tensor.dptr_, aux.dptr_,
      prefix_sum.data(), numel);
    // set the output shape forcefully and swap axis back
    mxnet::TShape out_shape(origin_shape);
    out_shape[actual_axis] = valid_num;
    const_cast<NDArray &>(outputs[0]).Init(out_shape);
    Tensor<cpu, 3, DType> output_tensor(outputs[0].data().dptr<DType>(),
        Shape3(temp_shape[1], valid_num, temp_shape[2]), stream);
    output_tensor = swapaxis<1, 0>(temp_tensor);
    // handle other optional outputs
    int output_flag = 0;
    if (param.return_index) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(mxnet::TShape(1, valid_num));
      dim_t* unique_indices = outputs[output_flag].data().dptr<dim_t>();
      // reuse boolean_mask kernel
      mxnet_op::Kernel<BooleanMaskForwardCPUKernel, cpu>::Launch(
        stream, temp_shape[0], unique_indices, perm.data(),
        prefix_sum.data(), 1);
    }
    if (param.return_inverse) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(mxnet::TShape(1, temp_shape[0]));
      dim_t* unique_inverse = outputs[output_flag].data().dptr<dim_t>();
      mxnet_op::Kernel<UniqueReturnInverseKernel, cpu>::Launch(
        stream, temp_shape[0], unique_inverse, prefix_sum.data(), perm.data());
    }
    if (param.return_counts) {
      output_flag += 1;
      std::vector<dim_t> idx(valid_num + 1);
      auto iter = idx.begin();
      for (dim_t i = 0; i < temp_shape[0]; ++i) {
        if (mask[i]) {
          *iter = i;
          ++iter;
        }
      }
      *iter = temp_shape[0];
      const_cast<NDArray &>(outputs[output_flag]).Init(mxnet::TShape(1, valid_num));
      dim_t* unique_counts = outputs[output_flag].data().dptr<dim_t>();
      mxnet_op::Kernel<UniqueReturnCountsKernel, cpu>::Launch(
          stream, valid_num, unique_counts, idx.data());
    }
  });
}

void NumpyUniqueCPUForward(const nnvm::NodeAttrs& attrs,
                           const OpContext &ctx,
                           const std::vector<NDArray> &inputs,
                           const std::vector<OpReqType> &req,
                           const std::vector<NDArray> &outputs) {
  CHECK_EQ(inputs.size(), 1U);
  CHECK(req[0] == kWriteTo || req[0] == kWriteInplace);
  using namespace mshadow;
  const NumpyUniqueParam& param = nnvm::get<NumpyUniqueParam>(attrs.parsed);
  if (inputs[0].shape().ndim() == 0) {
    CHECK(!param.axis.has_value() || param.axis.value() == -1 || param.axis.value() == 0)
      << "Axis can only be -1 or 0 for scalor tensor";
    Stream<cpu> *s = ctx.get_stream<cpu>();
    mxnet::TShape shape_1(1, 1);
    const_cast<NDArray &>(outputs[0]).Init(shape_1);
    MSHADOW_TYPE_SWITCH(outputs[0].dtype(), DType, {
      Tensor<cpu, 1, DType> unique_out =
          outputs[0].data().get_with_shape<cpu, 1, DType>(Shape1(1), s);
      ASSIGN_DISPATCH(unique_out, OpReqType::kWriteTo, inputs[0].data().dptr<DType>()[0]);
    });
    int output_flag = 0;
    if (param.return_index) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(shape_1);
      Tensor<cpu, 1, dim_t> outdata =
          outputs[output_flag].data().get_with_shape<cpu, 1, dim_t>(Shape1(1), s);
      ASSIGN_DISPATCH(outdata, OpReqType::kWriteTo, 0);
    }
    if (param.return_inverse) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(shape_1);
      Tensor<cpu, 1, dim_t> outdata =
          outputs[output_flag].data().get_with_shape<cpu, 1, dim_t>(Shape1(1), s);
      ASSIGN_DISPATCH(outdata, OpReqType::kWriteTo, 0);
    }
    if (param.return_counts) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(shape_1);
      Tensor<cpu, 1, dim_t> outdata =
          outputs[output_flag].data().get_with_shape<cpu, 1, dim_t>(Shape1(1), s);
      ASSIGN_DISPATCH(outdata, OpReqType::kWriteTo, 1);
    }
  } else if (inputs[0].shape().Size() == 0) {
    // If the input tensor is zero size, only a check on axis is needed
    if (param.axis.has_value()) {
      int axis = param.axis.value();
      if (axis < 0) axis += inputs[0].shape().ndim();
      CHECK(axis >= 0 && axis < inputs[0].shape().ndim())
        << "Axis must be within the range of input tensor's dimension";
    }
    // set the shapes of outputs
    mxnet::TShape shape_0(1, 0);
    const_cast<NDArray &>(outputs[0]).Init(shape_0);
    int output_flag = 0;
    if (param.return_index) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(shape_0);
    }
    if (param.return_inverse) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(shape_0);
    }
    if (param.return_counts) {
      output_flag += 1;
      const_cast<NDArray &>(outputs[output_flag]).Init(shape_0);
    }
  } else {
    if (!param.axis.has_value()) {
      NumpyUniqueCPUNoneAxisImpl(param, ctx, inputs, req, outputs);
    } else {
      NumpyUniqueCPUImpl(param, ctx, inputs, req, outputs);
    }
  }
}

DMLC_REGISTER_PARAMETER(NumpyUniqueParam);

NNVM_REGISTER_OP(_npi_unique)
.set_attr_parser(ParamParser<NumpyUniqueParam>)
.set_num_inputs(1)
.set_num_outputs([](const NodeAttrs& attrs) {
    const NumpyUniqueParam& param = nnvm::get<NumpyUniqueParam>(attrs.parsed);
    int output_num = 1;
    if (param.return_index) output_num += 1;
    if (param.return_inverse) output_num += 1;
    if (param.return_counts) output_num += 1;
    return output_num;
  })
.set_attr<nnvm::FListInputNames>("FListInputNames",
  [](const NodeAttrs& attrs) {
    return std::vector<std::string>{"data"};
  })
.set_attr<nnvm::FInferType>("FInferType", NumpyUniqueType)
.set_attr<FComputeEx>("FComputeEx<cpu>", NumpyUniqueCPUForward)
.set_attr<FInferStorageType>("FInferStorageType", NumpyUniqueStorageType)
.set_attr<FResourceRequest>("FResourceRequest",
  [](const NodeAttrs& attrs) {
    return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
  })
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
.add_argument("data", "NDArray-or-Symbol", "The input array")
.add_arguments(NumpyUniqueParam::__FIELDS__());

}  // namespace op
}  // namespace mxnet