// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#pragma once

#include <cub/config.cuh>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <cub/device/dispatch/dispatch_reduce.cuh>
#include <cub/iterator/arg_index_input_iterator.cuh>

#include <thrust/iterator/iterator_adaptor.h>

#include <cuda/__iterator/tabulate_output_iterator.h>
#include <cuda/std/__functional/identity.h>
#include <cuda/std/__utility/swap.h>
#include <cuda/std/limits>

#ifndef _CCCL_DOXYGEN_INVOKED // Do not document

CUB_NAMESPACE_BEGIN

namespace detail::reduce
{
template <typename GlobalAccumT, typename PromoteToGlobalOpT, typename GlobalReductionOpT, typename FinalResultOutIteratorT>
struct accumulating_transform_output_op
{
  bool first_partition;
  bool last_partition;

  // We use a double-buffer to make assignment idempotent (i.e., allow potential repeated assignment)
  GlobalAccumT* d_previous_aggregate;
  GlobalAccumT* d_aggregate_out;

  // Output iterator to which the final result of type `GlobalAccumT` across all partitions will be assigned
  FinalResultOutIteratorT d_out;

  // Unary promotion operator type that is used to transform a per-partition result to a global result
  PromoteToGlobalOpT promote_op;

  // Reduction operation
  GlobalReductionOpT reduce_op;

  template <typename IndexT, typename AccumT>
  _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(IndexT, AccumT per_partition_aggregate)
  {
    // Add this partitions aggregate to the global aggregate
    if (first_partition)
    {
      *d_aggregate_out = promote_op(per_partition_aggregate);
    }
    else
    {
      *d_aggregate_out = reduce_op(*d_previous_aggregate, promote_op(per_partition_aggregate));
    }

    // If this is the last partition, we write the global aggregate to the user-provided iterator
    if (last_partition)
    {
      *d_out = *d_aggregate_out;
    }
  }

  /**
   * This is a helper function that's invoked after a partition has been fully processed
   */
  template <typename GlobalOffsetT>
  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(GlobalOffsetT partition_size, bool next_partition_is_the_last)
  {
    promote_op.advance(partition_size);
    using ::cuda::std::swap;
    swap(d_previous_aggregate, d_aggregate_out);
    first_partition = false;
    last_partition  = next_partition_is_the_last;
  }
};

/**
 * Unary "promotion" operator type that is used to transform a per-partition result to a global result
 */
template <typename GlobalOffsetT>
struct local_to_global_op
{
  // The current partition's offset to be factored into this partition's index
  GlobalOffsetT current_partition_offset;

  /**
   * This helper function is invoked after a partition has been fully processed, in preparation for the next partition.
   */
  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(GlobalOffsetT partition_size)
  {
    current_partition_offset += partition_size;
  }

  /**
   * Unary operator called to transform the per-partition aggregate of a partition to a global aggregate type (i.e., one
   * that is used to reduce across partitions).
   */
  template <typename PerPartitionOffsetT, typename AccumT>
  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE KeyValuePair<GlobalOffsetT, AccumT>
  operator()(KeyValuePair<PerPartitionOffsetT, AccumT> partition_aggregate)
  {
    return KeyValuePair<GlobalOffsetT, AccumT>{
      current_partition_offset + static_cast<GlobalOffsetT>(partition_aggregate.key), partition_aggregate.value};
  }
};

template <typename ExtremumOutIteratorT, typename IndexOutIteratorT>
struct unzip_and_write_arg_extremum_op
{
  ExtremumOutIteratorT result_out_it;
  IndexOutIteratorT index_out_it;

  template <typename IndexT, typename KeyValuePairT>
  _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(IndexT, KeyValuePairT reduced_result)
  {
    *result_out_it = reduced_result.value;
    *index_out_it  = reduced_result.key;
  }
};

/******************************************************************************
 * Single-problem streaming reduction dispatch
 *****************************************************************************/

// Internal dispatch routine for computing a device-wide argument extremum, like `ArgMin` and `ArgMax`.
// Streaming, here, refers to the approach used for large number of items that are processed in multiple partitions.
//
// @tparam PerPartitionOffsetT
//   Offset type used as the index to access items within one partition, i.e., the offset type used within the kernel
//   template specialization
//
// @tparam InputIteratorT
//   Random-access input iterator type for reading input items @iterator
//
// @tparam OutputIteratorT
//   Output iterator type for writing the result of the (index, extremum)-key-value-pair
//
// @tparam GlobalOffsetT
//   Offset type used as the index to access items within the total input range, i.e., in the range [d_in, d_in +
// num_items)
//
// @tparam ReductionOpT
//   Binary reduction functor type having a member function that returns the selected extremum of two input items.
//   The streaming reduction requires two overloads, one used for selecting the extremum within one partition and one
//   for selecting the extremum across partitions.
//
// @tparam InitT
//   Initial value type
//
// @tparam PolicySelector
//   Selects the tuning policy
template <typename PerPartitionOffsetT,
          typename InputIteratorT,
          typename ExtremumOutIteratorT,
          typename IndexOutIteratorT,
          typename GlobalOffsetT,
          typename ReductionOpT,
          typename PolicySelector = policy_selector_from_types<
            KeyValuePair<PerPartitionOffsetT, non_void_value_t<ExtremumOutIteratorT, it_value_t<InputIteratorT>>>,
            PerPartitionOffsetT,
            ReductionOpT>>
#  if _CCCL_HAS_CONCEPTS()
  requires reduce_policy_selector<PolicySelector>
#  endif // _CCCL_HAS_CONCEPTS()
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_streaming_arg_reduce(
  void* d_temp_storage,
  size_t& temp_storage_bytes,
  InputIteratorT d_in,
  ExtremumOutIteratorT d_min_out,
  IndexOutIteratorT d_index_out,
  GlobalOffsetT num_items,
  ReductionOpT reduce_op,
  cudaStream_t stream,
  PolicySelector policy_selector = {})
{
  using input_value_t     = it_value_t<InputIteratorT>;
  using output_extremum_t = non_void_value_t<ExtremumOutIteratorT, input_value_t>;

  // Tabulate output iterator that unzips the result and writes it to the user-provided output iterators
  auto d_result_out = ::cuda::make_tabulate_output_iterator(
    detail::reduce::unzip_and_write_arg_extremum_op<ExtremumOutIteratorT, IndexOutIteratorT>{d_min_out, d_index_out});

  // Wrapped input iterator to produce index-value tuples, i.e., <PerPartitionOffsetT, InputT>-tuples
  // We make sure to offset the user-provided input iterator by the current partition's offset
  using arg_index_input_iterator_t = ArgIndexInputIterator<InputIteratorT, PerPartitionOffsetT, output_extremum_t>;

  // The output tuple type (i.e., extremum plus index tuples)
  using per_partition_accum_t = KeyValuePair<PerPartitionOffsetT, output_extremum_t>;
  using global_accum_t        = KeyValuePair<GlobalOffsetT, output_extremum_t>;

  // Unary promotion operator type that is used to transform a per-partition result to a global result
  // operator()(per_partition_accum_t) -> global_accum_t
  using local_to_global_op_t = local_to_global_op<GlobalOffsetT>;

  // The current partition's input iterator is an ArgIndex iterator that generates indices relative to the beginning
  // of the current partition, i.e., [0, partition_size) along with an OffsetIterator that offsets the user-provided
  // input iterator by the current partition's offset
  arg_index_input_iterator_t d_indexed_offset_in(d_in);

  // Transforms the per-partition result to a global result by adding the current partition's offset to the arg result
  // of a partition
  local_to_global_op_t local_to_global_op{GlobalOffsetT{0}};

  // Upper bound at which we want to cut the input into multiple partitions. Align to 4096 bytes for performance
  // reasons
  static constexpr PerPartitionOffsetT max_offset_size = ::cuda::std::numeric_limits<PerPartitionOffsetT>::max();
  static constexpr PerPartitionOffsetT max_partition_size =
    max_offset_size - (max_offset_size % PerPartitionOffsetT{4096});

  // Whether the given number of items fits into a single partition
  const bool is_single_partition =
    static_cast<GlobalOffsetT>(max_partition_size) >= static_cast<GlobalOffsetT>(num_items);

  // The largest partition size ever encountered
  const auto largest_partition_size =
    is_single_partition ? static_cast<PerPartitionOffsetT>(num_items) : max_partition_size;

  // Reduction operator type that enables accumulating per-partition results to a global reduction result
  using accumulating_transform_output_op_t =
    accumulating_transform_output_op<global_accum_t, local_to_global_op_t, ReductionOpT, decltype(d_result_out)>;
  auto accumulating_out_op = accumulating_transform_output_op_t{
    true, is_single_partition, nullptr, nullptr, d_result_out, local_to_global_op, reduce_op};

  // Initial value for empty problems, according to documented contract
  const auto empty_problem_extremum = static_cast<output_extremum_t>([] {
    if constexpr (::cuda::std::is_same_v<ReductionOpT, arg_min>
                  && ::cuda::std::numeric_limits<input_value_t>::is_specialized)
    {
      return ::cuda::std::numeric_limits<input_value_t>::max();
    }
    else if constexpr (::cuda::std::is_same_v<ReductionOpT, arg_max>
                       && ::cuda::std::numeric_limits<input_value_t>::is_specialized)
    {
      return ::cuda::std::numeric_limits<input_value_t>::lowest();
    }
    else
    {
      return input_value_t{};
    }
  }());
  auto initial_value = empty_problem_init_t<per_partition_accum_t>{{PerPartitionOffsetT{1}, empty_problem_extremum}};

  void* allocations[2]       = {nullptr, nullptr};
  size_t allocation_sizes[2] = {0, 2 * sizeof(global_accum_t)};

  // Query temporary storage requirements for per-partition reduction
  reduce::dispatch<per_partition_accum_t>(
    nullptr,
    allocation_sizes[0],
    d_indexed_offset_in,
    ::cuda::make_tabulate_output_iterator(accumulating_out_op),
    static_cast<PerPartitionOffsetT>(largest_partition_size),
    reduce_op,
    initial_value,
    stream,
    ::cuda::std::identity{},
    policy_selector);

  // Alias the temporary allocations from the single storage blob (or compute the necessary size of the blob)
  if (const auto error = CubDebug(alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
  {
    return error;
  }

  // Return if the caller is simply requesting the size of the storage allocation
  if (d_temp_storage == nullptr)
  {
    return cudaSuccess;
  }

  // Pointer to the double-buffer of global accumulators, which aggregate cross-partition results
  global_accum_t* const d_global_aggregates = static_cast<global_accum_t*>(allocations[1]);

  accumulating_out_op.d_previous_aggregate = d_global_aggregates;
  accumulating_out_op.d_aggregate_out      = d_global_aggregates + 1;

  for (GlobalOffsetT current_partition_offset = 0; current_partition_offset < static_cast<GlobalOffsetT>(num_items);
       current_partition_offset += static_cast<GlobalOffsetT>(max_partition_size))
  {
    const GlobalOffsetT remaining_items = (num_items - current_partition_offset);
    const GlobalOffsetT current_num_items =
      (remaining_items < max_partition_size) ? remaining_items : max_partition_size;

    d_indexed_offset_in = arg_index_input_iterator_t(d_in + current_partition_offset);

    if (const auto error = reduce::dispatch<per_partition_accum_t>(
          d_temp_storage,
          temp_storage_bytes,
          d_indexed_offset_in,
          ::cuda::make_tabulate_output_iterator(accumulating_out_op),
          static_cast<PerPartitionOffsetT>(current_num_items),
          reduce_op,
          initial_value,
          stream,
          ::cuda::std::identity{},
          policy_selector))
    {
      return error;
    }

    // Whether the next partition will be the last partition
    const bool next_partition_is_last =
      (remaining_items - current_num_items) <= static_cast<GlobalOffsetT>(max_partition_size);
    accumulating_out_op.advance(current_num_items, next_partition_is_last);
  }

  return cudaSuccess;
}
} // namespace detail::reduce
CUB_NAMESPACE_END

#endif // !_CCCL_DOXYGEN_INVOKED
