// Copyright (c) 2023 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.

#include "paddle/phi/kernels/sparse/unary_kernel.h"

#include "paddle/common/ddim.h"
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/slice_utils.h"

namespace phi::sparse {

template <typename T, typename Context>
void SliceCooCompute(const Context& dev_ctx,
                     const SparseCooTensor& x,
                     const std::vector<int64_t>& axes,
                     const std::vector<int64_t>& starts,
                     const std::vector<int64_t>& ends,
                     SparseCooTensor* out) {
  const DDim& x_dims = x.dims();

  // Step1: Infer output dims
  auto out_dims = funcs::GetSliceDims<int64_t>(
      x_dims, axes, starts, ends, nullptr, nullptr);

  // Step2: Get out_nnz (the number of non-zero elements in output)
  const int64_t x_nnz = x.nnz();
  int64_t out_nnz = 0;
  const auto* x_indices_data = x.indices().data<int64_t>();
  for (int64_t j = 0; j < x_nnz; ++j) {
    bool hit = true;
    for (size_t ii = 0; ii < axes.size(); ++ii) {
      auto item = x_indices_data[axes[ii] * x_nnz + j];
      if (!(starts[ii] <= item && item < ends[ii])) {
        hit = false;
        break;
      }
    }
    if (!hit) continue;
    out_nnz++;
  }

  // Step3: Get the values and indices of output
  auto sparse_dim = static_cast<int64_t>(x.sparse_dim());
  DenseTensor out_indices =
      Empty<int64_t, Context>(dev_ctx, {sparse_dim, out_nnz});
  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});

  auto* out_indices_data = out_indices.data<int64_t>();
  auto* out_values_data = out_values.data<T>();
  const auto* x_values_data = x.values().data<T>();
  int64_t index = 0;
  for (int64_t j = 0; j < x_nnz && index < out_nnz; ++j) {
    bool hit = true;
    for (size_t ii = 0; ii < axes.size(); ++ii) {
      auto item = x_indices_data[axes[ii] * x_nnz + j];
      if (!(starts[ii] <= item && item < ends[ii])) {
        hit = false;
        break;
      }
    }
    if (!hit) continue;
    // set value
    out_values_data[index] = x_values_data[j];
    // set coordinate
    for (int64_t i = 0; i < sparse_dim; ++i) {
      out_indices_data[i * out_nnz + index] = x_indices_data[i * x_nnz + j];
    }
    for (size_t ii = 0; ii < axes.size(); ++ii) {
      auto i = axes[ii];
      out_indices_data[i * out_nnz + index] -= starts[ii];
    }
    index++;
  }
  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
}

template <typename T, typename Context>
void SliceCooKernel(const Context& dev_ctx,
                    const SparseCooTensor& x,
                    const phi::IntArray& axes,
                    const phi::IntArray& starts,
                    const phi::IntArray& ends,
                    SparseCooTensor* out) {
  const DDim& x_dims = x.dims();
  std::vector<int64_t> axes_vec = axes.GetData();
  std::vector<int64_t> starts_vec = starts.GetData();
  std::vector<int64_t> ends_vec = ends.GetData();

  // Check and update attr
  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
      x_dims, &axes_vec, &starts_vec, &ends_vec);

  SliceCooCompute<T, Context>(dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
}

int64_t GetCsrNonZeroNumber(const SparseCsrTensor& x,
                            const int64_t x_crows_start,
                            const int64_t x_crows_end,
                            const int64_t min_col,
                            const int64_t max_col,
                            const int64_t x_cols_offset = 0) {
  const auto* x_crows_data = x.crows().data<int64_t>();
  const auto* x_cols_data = x.cols().data<int64_t>();
  int64_t out_nnz = 0;
  for (int64_t i = x_crows_start; i < x_crows_end; ++i) {
    int64_t st = x_crows_data[i] + x_cols_offset;
    int64_t ed = x_crows_data[i + 1] + x_cols_offset;
    for (int64_t jj = st; jj < ed; ++jj) {
      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
        out_nnz++;
      }
    }
  }
  return out_nnz;
}

template <typename T>
void GetCsrSubMatrix(const SparseCsrTensor& x,
                     const int64_t x_crows_start,
                     const int64_t x_crows_end,
                     const int64_t min_col,
                     const int64_t max_col,
                     DenseTensor* out_crows,
                     DenseTensor* out_cols,
                     DenseTensor* out_values,
                     const int64_t x_cols_offset = 0,
                     const int64_t out_crows_offset = 0,
                     const int64_t out_cols_offset = 0) {
  const auto* x_crows_data = x.crows().data<int64_t>();
  const auto* x_cols_data = x.cols().data<int64_t>();
  const auto* x_values_data = x.values().data<T>();

  auto* out_crows_data = out_crows->data<int64_t>();
  auto* out_cols_data = out_cols->data<int64_t>();
  auto* out_values_data = out_values->data<T>();
  out_crows_data[out_crows_offset] = 0;
  int64_t index = 0, out_n_rows = x_crows_end - x_crows_start;
  for (int i = 0; i < out_n_rows; ++i) {
    int64_t st = x_crows_data[x_crows_start + i] + x_cols_offset;
    int64_t ed = x_crows_data[x_crows_start + i + 1] + x_cols_offset;
    for (int64_t jj = st; jj < ed; ++jj) {
      if (x_cols_data[jj] >= min_col && x_cols_data[jj] < max_col) {
        out_cols_data[out_cols_offset + index] = x_cols_data[jj] - min_col;
        out_values_data[out_cols_offset + index] = x_values_data[jj];
        index++;
      }
    }
    out_crows_data[out_crows_offset + i + 1] = index;
  }
}

template <typename T, typename Context>
void SliceCsrTensor2D(const Context& dev_ctx,
                      const SparseCsrTensor& x,
                      const std::vector<int64_t>& axes,
                      const std::vector<int64_t>& starts,
                      const std::vector<int64_t>& ends,
                      const DDim& out_dims,
                      SparseCsrTensor* out) {
  // Step1: Get nnz of out
  int64_t out_nnz =
      GetCsrNonZeroNumber(x, starts[0], ends[0], starts[1], ends[1], 0);
  // Step2: Set out
  int64_t out_n_rows = ends[0] - starts[0];
  DenseTensor out_crows = Empty<int64_t, Context>(dev_ctx, {out_n_rows + 1});
  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});
  GetCsrSubMatrix<T>(x,
                     starts[0],
                     ends[0],
                     starts[1],
                     ends[1],
                     &out_crows,
                     &out_cols,
                     &out_values,
                     0,
                     0,
                     0);
  out->SetMember(out_crows, out_cols, out_values, out_dims);
}

template <typename T, typename Context>
void SliceCsrTensor3D(const Context& dev_ctx,
                      const SparseCsrTensor& x,
                      const std::vector<int64_t>& axes,
                      const std::vector<int64_t>& starts,
                      const std::vector<int64_t>& ends,
                      const DDim& out_dims,
                      SparseCsrTensor* out) {
  const auto* x_crows_data = x.crows().data<int64_t>();
  // Step1: Get nnz of out
  const int64_t x_dim0 = x.dims()[0], x_n_rows = x.dims()[1];
  int64_t x_cols_offset = 0, out_nnz = 0;
  // all_nnzs stores the nnz along with out_dim0, which will be used to set out.
  std::vector<int64_t> all_nnzs(ends[0] - starts[0]);
  for (int64_t i = 0; i < x_dim0; ++i) {
    if (i >= starts[0] && i < ends[0]) {  // slice dim 0
      int64_t x_crows_st = i * (x_n_rows + 1) + starts[1];
      int64_t x_crows_ed = i * (x_n_rows + 1) + ends[1];
      int64_t nnz = GetCsrNonZeroNumber(
          x, x_crows_st, x_crows_ed, starts[2], ends[2], x_cols_offset);
      out_nnz += nnz;
      all_nnzs[i - starts[0]] = nnz;
    }
    // get the start index in non_zero_cols_
    x_cols_offset += x_crows_data[(i + 1) * (x_n_rows + 1) - 1];
  }

  // Step2: Set out
  const int64_t out_dim0 = out_dims[0], out_n_rows = out_dims[1];
  DenseTensor out_crows =
      Empty<int64_t, Context>(dev_ctx, {out_dim0 * (out_n_rows + 1)});
  DenseTensor out_cols = Empty<int64_t, Context>(dev_ctx, {out_nnz});
  DenseTensor out_values = Empty<T, Context>(dev_ctx, {out_nnz});

  x_cols_offset = 0;
  int64_t out_crows_offset = 0, out_cols_offset = 0;
  for (int64_t i = 0; i < x_dim0; ++i) {
    if (i >= starts[0] && i < ends[0]) {  // slice dim 0
      int64_t x_crows_start = i * (x_n_rows + 1) + starts[1];
      int64_t x_crows_end = i * (x_n_rows + 1) + ends[1];
      GetCsrSubMatrix<T>(x,
                         x_crows_start,
                         x_crows_end,
                         starts[2],
                         ends[2],
                         &out_crows,
                         &out_cols,
                         &out_values,
                         x_cols_offset,
                         out_crows_offset,
                         out_cols_offset);
      out_crows_offset += (out_n_rows + 1);
      out_cols_offset += all_nnzs[i - starts[0]];
    }
    x_cols_offset += x_crows_data[(i + 1) * (x_n_rows + 1) - 1];
  }
  out->SetMember(out_crows, out_cols, out_values, out_dims);
}

template <typename T, typename Context>
void SliceCsrCompute(const Context& dev_ctx,
                     const SparseCsrTensor& x,
                     const std::vector<int64_t>& axes,
                     const std::vector<int64_t>& starts,
                     const std::vector<int64_t>& ends,
                     SparseCsrTensor* out) {
  const DDim& x_dims = x.dims();

  // Step1: Infer output dims
  auto out_dims = funcs::GetSliceDims<int64_t>(
      x_dims, axes, starts, ends, nullptr, nullptr);

  // Step2: Construct new axes, starts and ends.
  std::vector<int64_t> new_axes(3), new_starts(3), new_ends(3);
  funcs::ConstructNewSliceAttrs(
      x_dims, axes, starts, ends, &new_axes, &new_starts, &new_ends);

  // Step3: Slice csr tensor according to its dimension
  if (x_dims.size() == 2) {
    SliceCsrTensor2D<T, Context>(
        dev_ctx, x, new_axes, new_starts, new_ends, out_dims, out);
  } else if (x_dims.size() == 3) {
    SliceCsrTensor3D<T, Context>(
        dev_ctx, x, new_axes, new_starts, new_ends, out_dims, out);
  } else {
    // throw exception
    common::errors::InvalidArgument(
        "Slice for Sparse CSR Tensor only support 2-D or 3-D, but got %d-D.",
        x_dims.size());
  }
}

template <typename T, typename Context>
void SliceCsrKernel(const Context& dev_ctx,
                    const SparseCsrTensor& x,
                    const phi::IntArray& axes,
                    const phi::IntArray& starts,
                    const phi::IntArray& ends,
                    SparseCsrTensor* out) {
  const DDim& x_dims = x.dims();
  std::vector<int64_t> axes_vec = axes.GetData();
  std::vector<int64_t> starts_vec = starts.GetData();
  std::vector<int64_t> ends_vec = ends.GetData();

  // Check and update attr
  funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
      x_dims, &axes_vec, &starts_vec, &ends_vec);
  SliceCsrCompute<T, Context>(dev_ctx, x, axes_vec, starts_vec, ends_vec, out);
}
}  // namespace phi::sparse

PD_REGISTER_KERNEL(slice_coo,
                   CPU,
                   ALL_LAYOUT,
                   phi::sparse::SliceCooKernel,
                   float,
                   double,
                   int8_t,
                   uint8_t,
                   int16_t,
                   int,
                   int64_t,
                   bool) {}

PD_REGISTER_KERNEL(slice_csr,
                   CPU,
                   ALL_LAYOUT,
                   phi::sparse::SliceCsrKernel,
                   float,
                   double,
                   int8_t,
                   uint8_t,
                   int16_t,
                   int,
                   int64_t,
                   bool) {}