/*
 * Copyright (c) 2021-2024, NVIDIA CORPORATION.
 *
 * 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 <cub/block/block_load.cuh>
#include <cub/block/block_radix_sort.cuh>
#include <cub/block/block_store.cuh>
#include <cub/cub.cuh>

namespace raft {
namespace sparse {
namespace distance {
namespace detail {
/**
 * Load-balanced sparse-matrix-sparse-matrix multiplication (SPMM) kernel with
 * sparse-matrix-sparse-vector multiplication layout (SPMV).
 * This is intended to be scheduled n_chunks_b times for each row of a.
 * The steps are as follows:
 *
 * 1. Load row from A into dense vector in shared memory.
 *    This can be further chunked in the future if necessary to support larger
 *    column sizes.
 * 2. Threads of block all step through chunks of B in parallel.
 *    When a new row is encountered in row_indices_b, a segmented
 *    reduction is performed across the warps and then across the
 *    block and the final value written out to host memory.
 *
 * Reference: https://www.icl.utk.edu/files/publications/2020/icl-utk-1421-2020.pdf
 *
 * @tparam value_idx index type
 * @tparam value_t value type
 * @tparam tpb threads per block configured on launch
 * @tparam rev if this is true, the reduce/accumulate functions are only
 *         executed when A[col] == 0.0. when executed before/after !rev
 *         and A & B are reversed, this allows the full symmetric difference
 *         and intersection to be computed.
 * @tparam kv_t data type stored in shared mem cache
 * @tparam product_f reduce function type (semiring product() function).
 *                  accepts two arguments of value_t and returns a value_t
 * @tparam accum_f accumulation function type (semiring sum() function).
 *                 accepts two arguments of value_t and returns a value_t
 * @tparam write_f function to write value out. this should be mathematically
 *                 equivalent to the accumulate function but implemented as
 *                 an atomic operation on global memory. Accepts two arguments
 *                 of value_t* and value_t and updates the value given by the
 *                 pointer.
 * @param[in] indptrA column pointer array for A
 * @param[in] indicesA column indices array for A
 * @param[in] dataA data array for A
 * @param[in] rowsB coo row array for B
 * @param[in] indicesB column indices array for B
 * @param[in] dataB data array for B
 * @param[in] m number of rows in A
 * @param[in] n number of rows in B
 * @param[in] dim number of features
 * @param[in] nnz_b number of nonzeros in B
 * @param[out] out array of size m*n
 * @param[in] n_blocks_per_row number of blocks of B per row of A
 * @param[in] chunk_size number of nnz for B to use for each row of A
 * @param[in] buffer_size amount of smem to use for each row of A
 * @param[in] product_func semiring product() function
 * @param[in] accum_func semiring sum() function
 * @param[in] write_func atomic semiring sum() function
 */
template <typename strategy_t,
          typename indptr_it,
          typename value_idx,
          typename value_t,
          bool rev,
          int tpb,
          typename product_f,
          typename accum_f,
          typename write_f>
RAFT_KERNEL balanced_coo_generalized_spmv_kernel(strategy_t strategy,
                                                 indptr_it indptrA,
                                                 value_idx* indicesA,
                                                 value_t* dataA,
                                                 value_idx nnz_a,
                                                 value_idx* rowsB,
                                                 value_idx* indicesB,
                                                 value_t* dataB,
                                                 value_idx m,
                                                 value_idx n,
                                                 int dim,
                                                 value_idx nnz_b,
                                                 value_t* out,
                                                 int n_blocks_per_row,
                                                 int chunk_size,
                                                 value_idx b_ncols,
                                                 product_f product_func,
                                                 accum_f accum_func,
                                                 write_f write_func)
{
  typedef cub::WarpReduce<value_t> warp_reduce;

  value_idx cur_row_a        = indptrA.get_row_idx(n_blocks_per_row);
  value_idx cur_chunk_offset = blockIdx.x % n_blocks_per_row;

  // chunk starting offset
  value_idx ind_offset = cur_chunk_offset * chunk_size * tpb;
  // how many total cols will be processed by this block (should be <= chunk_size * n_threads)
  value_idx active_chunk_size = min(chunk_size * tpb, nnz_b - ind_offset);

  int tid     = threadIdx.x;
  int warp_id = tid / raft::warp_size();

  // compute id relative to current warp
  unsigned int lane_id = tid & (raft::warp_size() - 1);
  value_idx ind        = ind_offset + threadIdx.x;

  extern __shared__ char smem[];

  typename strategy_t::smem_type A                = (typename strategy_t::smem_type)(smem);
  typename warp_reduce::TempStorage* temp_storage = (typename warp_reduce::TempStorage*)(A + dim);

  auto inserter = strategy.init_insert(A, dim);

  __syncthreads();

  value_idx start_offset_a, stop_offset_a;
  bool first_a_chunk, last_a_chunk;
  indptrA.get_row_offsets(
    cur_row_a, start_offset_a, stop_offset_a, n_blocks_per_row, first_a_chunk, last_a_chunk);

  // Convert current row vector in A to dense
  for (int i = tid; i <= (stop_offset_a - start_offset_a); i += blockDim.x) {
    strategy.insert(inserter, indicesA[start_offset_a + i], dataA[start_offset_a + i]);
  }

  __syncthreads();

  auto finder = strategy.init_find(A, dim);

  if (cur_row_a > m || cur_chunk_offset > n_blocks_per_row) return;
  if (ind >= nnz_b) return;

  value_idx start_index_a = 0, stop_index_a = b_ncols - 1;
  indptrA.get_indices_boundary(indicesA,
                               cur_row_a,
                               start_offset_a,
                               stop_offset_a,
                               start_index_a,
                               stop_index_a,
                               first_a_chunk,
                               last_a_chunk);

  value_idx cur_row_b = -1;
  value_t c           = 0.0;

  auto warp_red = warp_reduce(*(temp_storage + warp_id));

  if (tid < active_chunk_size) {
    cur_row_b = rowsB[ind];

    auto index_b   = indicesB[ind];
    auto in_bounds = indptrA.check_indices_bounds(start_index_a, stop_index_a, index_b);

    if (in_bounds) {
      value_t a_col = strategy.find(finder, index_b);
      if (!rev || a_col == 0.0) { c = product_func(a_col, dataB[ind]); }
    }
  }

  // loop through chunks in parallel, reducing when a new row is
  // encountered by each thread
  for (int i = tid; i < active_chunk_size; i += blockDim.x) {
    value_idx ind_next   = ind + blockDim.x;
    value_idx next_row_b = -1;

    if (i + blockDim.x < active_chunk_size) next_row_b = rowsB[ind_next];

    bool diff_rows = next_row_b != cur_row_b;

    if (__any_sync(0xffffffff, diff_rows)) {
      // grab the threads currently participating in loops.
      // because any other threads should have returned already.
      unsigned int peer_group = __match_any_sync(0xffffffff, cur_row_b);
      bool is_leader          = get_lowest_peer(peer_group) == lane_id;
      value_t v               = warp_red.HeadSegmentedReduce(c, is_leader, accum_func);

      // thread with lowest lane id among peers writes out
      if (is_leader && v != 0.0) {
        // this conditional should be uniform, since rev is constant
        size_t idx = !rev ? (size_t)cur_row_a * n + cur_row_b : (size_t)cur_row_b * m + cur_row_a;
        write_func(out + idx, v);
      }

      c = 0.0;
    }

    if (next_row_b != -1) {
      ind = ind_next;

      auto index_b   = indicesB[ind];
      auto in_bounds = indptrA.check_indices_bounds(start_index_a, stop_index_a, index_b);
      if (in_bounds) {
        value_t a_col = strategy.find(finder, index_b);

        if (!rev || a_col == 0.0) { c = accum_func(c, product_func(a_col, dataB[ind])); }
      }

      cur_row_b = next_row_b;
    }
  }
}

}  // namespace detail
}  // namespace distance
}  // namespace sparse
}  // namespace raft
