// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// 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 <thrust/adjacent_difference.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/functional.h>
#include <thrust/scatter.h>
#include <thrust/sequence.h>
#include <thrust/unique.h>

#include <iostream>
#include <vector>

#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/funcs/unique_functor.h"

namespace phi {

// The core logic of computing Unique Consecutive for a flattened Tensor
template <typename Context,
          typename InT,
          typename IndexT,
          typename equal_T,
          typename not_equal_T>
static void UniqueConsecutiveFlattenedCUDATensor(const Context& dev_ctx,
                                                 const DenseTensor& in,
                                                 DenseTensor* out,
                                                 bool return_inverse,
                                                 bool return_counts,
                                                 equal_T equal,
                                                 not_equal_T not_equal,
                                                 int64_t num_input,
                                                 DenseTensor* inverse,
                                                 DenseTensor* counts) {
  // 0. Preparation
  DenseTensor in_hat;
  Copy(dev_ctx, in, dev_ctx.GetPlace(), false, &in_hat);
  auto in_data_hat = dev_ctx.template Alloc<InT>(&in_hat);

  DenseTensor sorted_indices;
  sorted_indices.Resize({num_input});
  auto sorted_indices_data = dev_ctx.template Alloc<IndexT>(&sorted_indices);
  thrust::sequence(
      thrust::device, sorted_indices_data, sorted_indices_data + num_input);
  // 1. Calculate op result: 'out'
  DenseTensor range;
  range.Resize({num_input + 1});
  auto range_data_ptr = dev_ctx.template Alloc<IndexT>(&range);
  thrust::sequence(
      thrust::device, range_data_ptr, range_data_ptr + num_input + 1);
  Copy(dev_ctx, in_hat, dev_ctx.GetPlace(), false, out);
  int num_out;
  auto out_data = dev_ctx.template Alloc<InT>(out);
  num_out =
      thrust::unique_by_key(
          thrust::device, out_data, out_data + num_input, range_data_ptr, equal)
          .first -
      out_data;
  out->Resize({num_out});

  // 2. Calculate inverse index: 'inverse'
  if (return_inverse) {
    inverse->Resize({num_input});
    auto inverse_data = dev_ctx.template Alloc<IndexT>(inverse);
    DenseTensor inv_loc;
    inv_loc.Resize({num_input});
    auto inv_loc_data_ptr = dev_ctx.template Alloc<IndexT>(&inv_loc);
    thrust::adjacent_difference(thrust::device,
                                in_data_hat,
                                in_data_hat + num_input,
                                inv_loc_data_ptr,
                                not_equal);
    thrust::device_ptr<IndexT> inv_loc_data_dev(inv_loc_data_ptr);
    inv_loc_data_dev[0] = 0;  // without device_ptr, segmentation fault
    thrust::inclusive_scan(thrust::device,
                           inv_loc_data_ptr,
                           inv_loc_data_ptr + num_input,
                           inv_loc_data_ptr);
    thrust::scatter(thrust::device,
                    inv_loc_data_ptr,
                    inv_loc_data_ptr + num_input,
                    sorted_indices_data,
                    inverse_data);
  }
  // 3. Calculate 'counts'
  if (return_counts) {
    counts->Resize({num_out});
    auto count_data = dev_ctx.template Alloc<IndexT>(counts);
    // init 'count_data' as 0
    thrust::fill(thrust::device, count_data, count_data + num_out, 0);
    thrust::device_ptr<IndexT> range_data_ptr_dev(range_data_ptr);
    range_data_ptr_dev[num_out] = num_input;
    thrust::adjacent_difference(thrust::device,
                                range_data_ptr + 1,
                                range_data_ptr + num_out + 1,
                                count_data);
  }
}

// functor for processing a flattened Tensor
template <typename Context, typename InT>
struct UniqueConsecutiveFlattenedCUDAFunctor {
  const Context& dev_ctx_;
  const DenseTensor& in_;
  DenseTensor* out_;
  const bool return_inverse_;
  const bool return_counts_;
  DenseTensor* inverse_;
  DenseTensor* count_;

  UniqueConsecutiveFlattenedCUDAFunctor(const Context& dev_ctx,
                                        const DenseTensor& in,
                                        DenseTensor* out,
                                        bool return_inverse,
                                        bool return_counts,
                                        DenseTensor* inverse,
                                        DenseTensor* count)
      : dev_ctx_(dev_ctx),
        in_(in),
        out_(out),
        return_inverse_(return_inverse),
        return_counts_(return_counts),
        inverse_(inverse),
        count_(count) {}

  template <typename IndexT>
  void apply() const {
    UniqueConsecutiveFlattenedCUDATensor<Context, InT, IndexT>(
        dev_ctx_,
        in_,
        out_,
        return_inverse_,
        return_counts_,
        thrust::equal_to<InT>(),
        thrust::not_equal_to<InT>(),
        in_.numel(),
        inverse_,
        count_);
  }
};

// The logic of compute unique with axis required, it's a little different
// from above function
template <typename Context,
          typename InT,
          typename IndexT,
          typename equal_T,
          typename not_equal_T>
static void ComputeUniqueConsecutiveDims(const Context& dev_ctx,
                                         DenseTensor* sorted_indices,
                                         IndexT* sorted_indices_data,
                                         DenseTensor* out,
                                         bool return_inverse,
                                         bool return_counts,
                                         equal_T equal,
                                         not_equal_T not_equal,
                                         int64_t row,
                                         DenseTensor* inverse,
                                         DenseTensor* counts) {
  // 1. inverse indices: 'inverse'
  DenseTensor tmp;
  if (!inverse) {
    inverse = &tmp;
  }

  inverse->Resize({row});
  auto inverse_data = dev_ctx.template Alloc<IndexT>(inverse);
  DenseTensor inv_loc;
  inv_loc.Resize({row});
  auto inv_loc_data_ptr = dev_ctx.template Alloc<IndexT>(&inv_loc);
  thrust::adjacent_difference(thrust::device,
                              sorted_indices_data,
                              sorted_indices_data + row,
                              inv_loc_data_ptr,
                              not_equal);
  thrust::device_ptr<IndexT> inv_loc_data_dev(inv_loc_data_ptr);
  inv_loc_data_dev[0] = 0;
  thrust::inclusive_scan(thrust::device,
                         inv_loc_data_ptr,
                         inv_loc_data_ptr + row,
                         inv_loc_data_ptr);
  thrust::scatter(thrust::device,
                  inv_loc_data_ptr,
                  inv_loc_data_ptr + row,
                  sorted_indices_data,
                  inverse_data);

  // 2. sorted indices
  DenseTensor range;
  range.Resize({row + 1});
  auto range_data_ptr = dev_ctx.template Alloc<IndexT>(&range);
  thrust::sequence(thrust::device, range_data_ptr, range_data_ptr + row + 1);
  int num_out;
  num_out = thrust::unique_by_key(thrust::device,
                                  sorted_indices_data,
                                  sorted_indices_data + row,
                                  range_data_ptr,
                                  equal)
                .first -
            sorted_indices_data;
  thrust::device_ptr<IndexT> range_data_ptr_dev(range_data_ptr);
  range_data_ptr_dev[num_out] = row;
  sorted_indices->Resize({num_out});

  // 3. counts: 'counts'
  if (return_counts) {
    counts->Resize({num_out});
    auto count_data = dev_ctx.template Alloc<IndexT>(counts);
    thrust::fill(thrust::device, count_data, count_data + row, 0);
    thrust::adjacent_difference(thrust::device,
                                range_data_ptr + 1,
                                range_data_ptr + row + 1,
                                count_data);
  }
}

// Binary function 'equal_to'
template <typename InT>
struct BinaryEqual {
  int64_t col;
  const InT* in_trans_data;

  BinaryEqual(int64_t _col, const InT* _in_trans_data)
      : col(_col), in_trans_data(_in_trans_data) {}

  __device__ bool operator()(int64_t a, int64_t b) const {
    for (int64_t i = 0; i < col; ++i) {
      InT lhs = in_trans_data[i + a * col];
      InT rhs = in_trans_data[i + b * col];
      if (lhs != rhs) {
        return false;
      }
    }
    return true;
  }
};

// Binary function 'not_equal_to'
template <typename InT>
struct BinaryNotEqual {
  int64_t col;
  const InT* in_trans_data;

  BinaryNotEqual(int64_t _col, const InT* _in_trans_data)
      : col(_col), in_trans_data(_in_trans_data) {}

  __device__ bool operator()(int64_t a, int64_t b) const {
    for (int64_t i = 0; i < col; ++i) {
      InT lhs = in_trans_data[i + a * col];
      InT rhs = in_trans_data[i + b * col];
      if (lhs != rhs) {
        return true;
      }
    }
    return false;
  }
};

// index_select() function for Tensor
template <typename Context, typename InT, typename IndexT>
void IndexSelect(const Context& dev_ctx,
                 const DenseTensor& input,
                 const DenseTensor& index,
                 DenseTensor* output,
                 int dim) {
  auto input_dim = input.dims();
  auto input_dim_size = input_dim.size();
  auto output_dim = output->dims();

  auto slice_size = 1;
  for (auto i = dim + 1; i < input_dim_size; i++) {
    slice_size *= input_dim[i];
  }

  auto input_width = slice_size * input_dim[dim];
  auto output_width = slice_size * output_dim[dim];

  auto outer_nums = 1;
  for (auto i = 0; i < dim; i++) {
    outer_nums *= input_dim[i];
  }

  auto index_size = index.dims()[0];

  std::vector<InT> input_vec;
  std::vector<IndexT> index_vec;
  TensorToVector(input, dev_ctx, &input_vec);
  TensorToVector(index, dev_ctx, &index_vec);
  std::vector<InT> out_vec(output->numel());

  for (int i = 0; i < index_size; i++) {
    PADDLE_ENFORCE_GE(
        index_vec[i],
        -input_dim[dim],
        common::errors::InvalidArgument(
            "Variable value (index) of OP(index_select) "
            "expected >= %ld and < %ld, but got %ld. Please check input "
            "value.",
            -input_dim[dim],
            input_dim[dim],
            index_vec[i]));
    PADDLE_ENFORCE_LT(
        index_vec[i],
        input_dim[dim],
        common::errors::InvalidArgument(
            "Variable value (index) of OP(index_select) "
            "expected >= %ld and < %ld, but got %ld. Please check input "
            "value.",
            -input_dim[dim],
            input_dim[dim],
            index_vec[i]));
  }

  for (int64_t i = 0; i < outer_nums; i++) {
    int64_t input_start_offset = i * input_width;
    int64_t output_start_offset = i * output_width;

    for (int64_t j = 0; j < index_size; j++) {
      IndexT index_value = index_vec[j];
      if (index_value < 0) {
        index_value += input_dim[dim];
      }
      for (int64_t k = 0; k < slice_size; k++) {
        out_vec[output_start_offset + j * slice_size + k] =
            input_vec[input_start_offset + index_value * slice_size + k];
      }
    }
  }
  dev_ctx.template Alloc<InT>(output);
  TensorFromVector(out_vec, dev_ctx, output);
  output->Resize(output_dim);
}

// Calculate unique consecutive when 'axis' is set
template <typename Context, typename InT, typename IndexT>
static void UniqueConsecutiveDimsCUDATensor(const Context& dev_ctx,
                                            const DenseTensor& in,
                                            DenseTensor* out,
                                            bool return_inverse,
                                            bool return_counts,
                                            int axis,
                                            DenseTensor* inverse,
                                            DenseTensor* counts) {
  // 1. Transpose & reshape
  // Transpose tensor: eg. axis=1, [dim0, dim1, dim2] -> [dim1, dim0, dim2]
  std::vector<int> permute(in.dims().size());
  std::iota(permute.begin(), permute.end(), 0);
  permute[axis] = 0;
  permute[0] = axis;
  std::vector<int64_t> in_trans_dims_vec(vectorize(in.dims()));
  in_trans_dims_vec[axis] = in.dims()[0];
  in_trans_dims_vec[0] = in.dims()[axis];
  DenseTensor in_trans;
  DDim in_trans_dims = make_ddim(in_trans_dims_vec);
  in_trans.Resize(in_trans_dims);
  dev_ctx.template Alloc<InT>(&in_trans);
  funcs::TransCompute<Context, InT>(in.dims().size(),  // num of dims
                                    dev_ctx,           // device
                                    in,                // original Tensor
                                    &in_trans,         // Tensor after reshape
                                    permute);          // index of axis

  // Reshape tensor: eg. [dim1, dim0, dim2] -> [dim1, dim0*dim2]
  DDim in_trans_flat_dims = common::flatten_to_2d(in_trans_dims, 1);
  in_trans.Resize(in_trans_flat_dims);

  // now 'in_trans' is 2D
  int64_t col = in_trans.dims()[1];
  int64_t row = in_trans.dims()[0];
  const InT* in_trans_data = in_trans.data<InT>();

  DenseTensor sorted_indices;
  sorted_indices.Resize({row});
  auto sorted_indices_data = dev_ctx.template Alloc<IndexT>(&sorted_indices);

  // 2. Calculate 'inverse', 'counts'
  // Init index
  thrust::sequence(
      thrust::device, sorted_indices_data, sorted_indices_data + row);
  ComputeUniqueConsecutiveDims<Context, InT, IndexT>(
      dev_ctx,
      &sorted_indices,
      sorted_indices_data,
      out,
      return_inverse,
      return_counts,
      BinaryEqual<InT>(col, in_trans_data),
      BinaryNotEqual<InT>(col, in_trans_data),
      row,
      inverse,
      counts);

  // 3. Select indices and reshape back to get 'out'
  DenseTensor out_trans;
  std::vector<int64_t> out_trans_dims_vec = in_trans_dims_vec;
  out_trans_dims_vec[0] = sorted_indices.numel();
  out_trans.Resize(out_trans_dims_vec);
  dev_ctx.template Alloc<InT>(&out_trans);

  IndexSelect<Context, InT, IndexT>(
      dev_ctx, in_trans, sorted_indices, &out_trans, 0);

  std::swap(out_trans_dims_vec[0], out_trans_dims_vec[axis]);
  out->Resize(out_trans_dims_vec);
  dev_ctx.template Alloc<InT>(out);
  std::vector<DenseTensor> out_trans_unbind = funcs::Unbind(out_trans);
  funcs::ConcatFunctor<Context, InT> concat_functor;
  concat_functor(dev_ctx, out_trans_unbind, 0, &out_trans);
  funcs::TransCompute<Context, InT>(
      out_trans.dims().size(), dev_ctx, out_trans, out, permute);
}

// functor for processing a multi-dimensional Tensor
template <typename Context, typename InT>
struct UniqueConsecutiveDimsCUDAFunctor {
  const Context& dev_ctx_;
  const DenseTensor& in_;
  DenseTensor* out_;
  const int axis_;
  const bool return_inverse_;
  const bool return_counts_;
  DenseTensor* inverse_;
  DenseTensor* count_;

  UniqueConsecutiveDimsCUDAFunctor(const Context& dev_ctx,
                                   const DenseTensor& in,
                                   DenseTensor* out,
                                   const int axis,
                                   bool return_inverse,
                                   bool return_counts,
                                   DenseTensor* inverse,
                                   DenseTensor* count)
      : dev_ctx_(dev_ctx),
        in_(in),
        out_(out),
        axis_(axis),
        return_inverse_(return_inverse),
        return_counts_(return_counts),
        inverse_(inverse),
        count_(count) {}

  template <typename IndexT>
  void apply() const {
    UniqueConsecutiveDimsCUDATensor<Context, InT, IndexT>(dev_ctx_,
                                                          in_,
                                                          out_,
                                                          return_inverse_,
                                                          return_counts_,
                                                          axis_,
                                                          inverse_,
                                                          count_);
  }
};

}  // namespace phi