// SPDX-FileCopyrightText: Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause

#pragma once

#include <thrust/detail/config.h>

#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

#if _CCCL_CUDA_COMPILATION()

#  include <thrust/system/cuda/config.h>

#  include <cub/device/dispatch/dispatch_select_if.cuh>
#  include <cub/util_device.cuh>
#  include <cub/util_math.cuh>

#  include <thrust/detail/raw_reference_cast.h>
#  include <thrust/detail/temporary_array.h>
#  include <thrust/partition.h>
#  include <thrust/system/cuda/detail/cdp_dispatch.h>
#  include <thrust/system/cuda/detail/execution_policy.h>
#  include <thrust/system/cuda/detail/find.h>
#  include <thrust/system/cuda/detail/reverse.h>
#  include <thrust/system/cuda/detail/uninitialized_copy.h>
#  include <thrust/system/cuda/detail/util.h>

#  include <cuda/__iterator/zip_iterator.h>
#  include <cuda/std/__iterator/distance.h>
#  include <cuda/std/__utility/pair.h>
#  include <cuda/std/cstdint>
#  include <cuda/std/tuple>

THRUST_NAMESPACE_BEGIN
namespace cuda_cub
{
namespace detail
{
template <typename Derived, typename InputIt, typename StencilIt, typename OutputIt, typename Predicate, typename OffsetT>
[[nodiscard]] cudaError_t THRUST_RUNTIME_FUNCTION dispatch_partition(
  execution_policy<Derived>& policy,
  void* d_temp_storage,
  size_t& temp_storage_bytes,
  InputIt first,
  StencilIt stencil,
  OutputIt output,
  Predicate predicate,
  OffsetT num_items,
  std::size_t& num_selected)
{
  // FIXME(bgruber): the call to the dispatch function should be replaced by a public CUB API, but there is currently no
  // exposure of a partition with stencil and predicate (a `DevicePartition:FlaggedIf`)

  using equality_op_t = cub::NullType;

  cudaStream_t stream = cuda_cub::stream(policy);

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

  // Query algorithm memory requirements
  cudaError_t status = cub::detail::select::dispatch<cub::SelectImpl::Partition>(
    nullptr,
    allocation_sizes[0],
    first,
    stencil,
    output,
    static_cast<OffsetT*>(nullptr),
    predicate,
    equality_op_t{},
    num_items,
    stream);
  _CUDA_CUB_RET_IF_FAIL(status);

  status = cub::detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes);
  _CUDA_CUB_RET_IF_FAIL(status);

  // Return if we're only querying temporary storage requirements
  if (d_temp_storage == nullptr)
  {
    return status;
  }

  // Return for empty problems
  if (num_items == 0)
  {
    num_selected = 0;
    return status;
  }

  // Memory allocation for the number of selected output items
  OffsetT* d_num_selected_out = thrust::detail::aligned_reinterpret_cast<OffsetT*>(allocations[1]);

  // Run algorithm
  status = cub::detail::select::dispatch<cub::SelectImpl::Partition>(
    allocations[0],
    allocation_sizes[0],
    first,
    stencil,
    output,
    d_num_selected_out,
    predicate,
    equality_op_t{},
    num_items,
    stream);
  _CUDA_CUB_RET_IF_FAIL(status);

  // Get number of selected items
  status = cuda_cub::synchronize(policy);
  _CUDA_CUB_RET_IF_FAIL(status);
  num_selected = static_cast<std::size_t>(get_value(policy, d_num_selected_out));

  return status;
}

template <typename Derived, typename InputIt, typename StencilIt, typename OutputIt, typename Predicate>
THRUST_RUNTIME_FUNCTION std::size_t partition(
  execution_policy<Derived>& policy,
  InputIt first,
  InputIt last,
  StencilIt stencil,
  OutputIt output,
  Predicate predicate)
{
  const auto num_items = ::cuda::std::distance(first, last);
  std::size_t num_selected{};
  cudaError_t status        = cudaSuccess;
  size_t temp_storage_bytes = 0;

  // Query temporary storage requirements
  THRUST_INDEX_TYPE_DISPATCH(
    status,
    dispatch_partition,
    num_items,
    (policy, nullptr, temp_storage_bytes, first, stencil, output, predicate, num_items_fixed, num_selected));
  cuda_cub::throw_on_error(status, "partition failed on 1st step");

  // Allocate temporary storage.
  thrust::detail::temporary_array<std::uint8_t, Derived> tmp(policy, temp_storage_bytes);
  void* temp_storage = static_cast<void*>(tmp.data().get());

  // Run algorithm
  THRUST_INDEX_TYPE_DISPATCH(
    status,
    dispatch_partition,
    num_items,
    (policy, temp_storage, temp_storage_bytes, first, stencil, output, predicate, num_items_fixed, num_selected));
  cuda_cub::throw_on_error(status, "partition failed on 2nd step");

  return num_selected;
}

template <typename Derived,
          typename InputIt,
          typename StencilIt,
          typename SelectedOutIt,
          typename RejectedOutIt,
          typename Predicate>
THRUST_RUNTIME_FUNCTION ::cuda::std::pair<SelectedOutIt, RejectedOutIt> stable_partition_copy(
  execution_policy<Derived>& policy,
  InputIt first,
  InputIt last,
  StencilIt stencil,
  SelectedOutIt selected_result,
  RejectedOutIt rejected_result,
  Predicate predicate)
{
  if (::cuda::std::distance(first, last) <= 0)
  {
    return ::cuda::std::make_pair(selected_result, rejected_result);
  }

  using output_it_wrapper_t = cub::detail::select::partition_distinct_output_t<SelectedOutIt, RejectedOutIt>;
  std::size_t num_items     = static_cast<std::size_t>(::cuda::std::distance(first, last));
  std::size_t num_selected =
    partition(policy, first, last, stencil, output_it_wrapper_t{selected_result, rejected_result}, predicate);
  return ::cuda::std::make_pair(selected_result + num_selected, rejected_result + num_items - num_selected);
}

template <typename Derived, typename InputIt, typename StencilIt, typename Predicate>
THRUST_RUNTIME_FUNCTION InputIt inplace_partition(
  execution_policy<Derived>& policy, InputIt first, InputIt last, StencilIt stencil, Predicate predicate)
{
  if (::cuda::std::distance(first, last) <= 0)
  {
    return first;
  }

  // Element type of the input iterator
  using value_t         = thrust::detail::it_value_t<InputIt>;
  std::size_t num_items = static_cast<std::size_t>(::cuda::std::distance(first, last));

  // Allocate temporary storage, which will serve as the input to the partition
  thrust::detail::temporary_array<value_t, Derived> tmp(policy, num_items);
  cuda_cub::uninitialized_copy(policy, first, last, tmp.begin());

  // Partition input from temporary storage to the user-provided range [`first`, `last`)
  std::size_t num_selected =
    partition(policy, tmp.data().get(), tmp.data().get() + num_items, stencil, first, predicate);
  return first + num_selected;
}
} // namespace detail

//-------------------------
// Thrust API entry points
//-------------------------

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class InputIt, class StencilIt, class SelectedOutIt, class RejectedOutIt, class Predicate>
::cuda::std::pair<SelectedOutIt, RejectedOutIt> _CCCL_HOST_DEVICE partition_copy(
  execution_policy<Derived>& policy,
  InputIt first,
  InputIt last,
  StencilIt stencil,
  SelectedOutIt selected_result,
  RejectedOutIt rejected_result,
  Predicate predicate)
{
  auto ret = ::cuda::std::make_pair(selected_result, rejected_result);
  THRUST_CDP_DISPATCH(
    (ret = detail::stable_partition_copy(policy, first, last, stencil, selected_result, rejected_result, predicate);),
    (ret = thrust::partition_copy(
       cvt_to_seq(derived_cast(policy)), first, last, stencil, selected_result, rejected_result, predicate);));
  return ret;
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class InputIt, class SelectedOutIt, class RejectedOutIt, class Predicate>
::cuda::std::pair<SelectedOutIt, RejectedOutIt> _CCCL_HOST_DEVICE partition_copy(
  execution_policy<Derived>& policy,
  InputIt first,
  InputIt last,
  SelectedOutIt selected_result,
  RejectedOutIt rejected_result,
  Predicate predicate)
{
  auto ret = ::cuda::std::make_pair(selected_result, rejected_result);
  THRUST_CDP_DISPATCH(
    (ret = detail::stable_partition_copy(
       policy, first, last, static_cast<cub::NullType*>(nullptr), selected_result, rejected_result, predicate);),
    (ret = thrust::partition_copy(
       cvt_to_seq(derived_cast(policy)), first, last, selected_result, rejected_result, predicate);));
  return ret;
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class InputIt, class StencilIt, class SelectedOutIt, class RejectedOutIt, class Predicate>
::cuda::std::pair<SelectedOutIt, RejectedOutIt> _CCCL_HOST_DEVICE stable_partition_copy(
  execution_policy<Derived>& policy,
  InputIt first,
  InputIt last,
  StencilIt stencil,
  SelectedOutIt selected_result,
  RejectedOutIt rejected_result,
  Predicate predicate)
{
  auto ret = ::cuda::std::make_pair(selected_result, rejected_result);
  THRUST_CDP_DISPATCH(
    (ret = detail::stable_partition_copy(policy, first, last, stencil, selected_result, rejected_result, predicate);),
    (ret = thrust::stable_partition_copy(
       cvt_to_seq(derived_cast(policy)), first, last, stencil, selected_result, rejected_result, predicate);));
  return ret;
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class InputIt, class SelectedOutIt, class RejectedOutIt, class Predicate>
::cuda::std::pair<SelectedOutIt, RejectedOutIt> _CCCL_HOST_DEVICE stable_partition_copy(
  execution_policy<Derived>& policy,
  InputIt first,
  InputIt last,
  SelectedOutIt selected_result,
  RejectedOutIt rejected_result,
  Predicate predicate)
{
  auto ret = ::cuda::std::make_pair(selected_result, rejected_result);
  THRUST_CDP_DISPATCH(
    (ret = detail::stable_partition_copy(
       policy, first, last, static_cast<cub::NullType*>(nullptr), selected_result, rejected_result, predicate);),
    (ret = thrust::stable_partition_copy(
       cvt_to_seq(derived_cast(policy)), first, last, selected_result, rejected_result, predicate);));
  return ret;
}

/// inplace

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class Iterator, class StencilIt, class Predicate>
Iterator _CCCL_HOST_DEVICE
partition(execution_policy<Derived>& policy, Iterator first, Iterator last, StencilIt stencil, Predicate predicate)
{
  THRUST_CDP_DISPATCH((last = detail::inplace_partition(policy, first, last, stencil, predicate);),
                      (last = thrust::partition(cvt_to_seq(derived_cast(policy)), first, last, stencil, predicate);));
  return last;
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class Iterator, class Predicate>
Iterator _CCCL_HOST_DEVICE
partition(execution_policy<Derived>& policy, Iterator first, Iterator last, Predicate predicate)
{
  THRUST_CDP_DISPATCH(
    (last = detail::inplace_partition(policy, first, last, static_cast<cub::NullType*>(nullptr), predicate);),
    (last = thrust::partition(cvt_to_seq(derived_cast(policy)), first, last, predicate);));
  return last;
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class Iterator, class StencilIt, class Predicate>
Iterator _CCCL_HOST_DEVICE stable_partition(
  execution_policy<Derived>& policy, Iterator first, Iterator last, StencilIt stencil, Predicate predicate)
{
  auto ret = last;
  THRUST_CDP_DISPATCH(
    (ret = detail::inplace_partition(policy, first, last, stencil, predicate);

     /* partition returns rejected values in reverse order
       so reverse the rejected elements to make it stable */
     cuda_cub::reverse(policy, ret, last);),
    (ret = thrust::stable_partition(cvt_to_seq(derived_cast(policy)), first, last, stencil, predicate);));
  return ret;
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class Iterator, class Predicate>
Iterator _CCCL_HOST_DEVICE
stable_partition(execution_policy<Derived>& policy, Iterator first, Iterator last, Predicate predicate)
{
  auto ret = last;
  THRUST_CDP_DISPATCH(
    (ret = detail::inplace_partition(policy, first, last, static_cast<cub::NullType*>(nullptr), predicate);

     /* partition returns rejected values in reverse order
      so reverse the rejected elements to make it stable */
     cuda_cub::reverse(policy, ret, last);),
    (ret = thrust::stable_partition(cvt_to_seq(derived_cast(policy)), first, last, predicate);));
  return ret;
}

// Functor for the single-pass is_partitioned check.
// Returns true for an adjacent pair (a[i], a[i+1]) where pred(a[i]) is false
// and pred(a[i+1]) is true — i.e., a "false → true" transition that violates
// the partitioning invariant.
template <class Predicate>
struct __is_partitioned_fn
{
  Predicate pred_;

  // Not const-qualified: a const operator() would propagate const onto pred_
  // and reject predicates whose own operator() is non-const (Thrust permits these).
  template <class Tuple>
  [[nodiscard]] _CCCL_HOST_DEVICE bool operator()(const Tuple& tuple)
  {
    const bool lhs = pred_(thrust::raw_reference_cast(::cuda::std::get<0>(tuple)));
    const bool rhs = pred_(thrust::raw_reference_cast(::cuda::std::get<1>(tuple)));
    return !lhs && rhs;
  }
};

// Single-pass implementation: zip adjacent elements and find any "false → true"
// transition. Two-pass (find_if_not + find_if) required two kernel launches;
// this approach uses one find_if over (a[i], a[i+1]) pairs, cutting kernel
// launch overhead roughly in half for typical inputs.
// See: https://github.com/NVIDIA/cccl/issues/8085
template <class Derived, class ItemsIt, class Predicate>
bool _CCCL_HOST_DEVICE
is_partitioned(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last, Predicate predicate)
{
  if (first == last)
  {
    return true;
  }
  // Build a range of adjacent pairs: (a[0],a[1]), (a[1],a[2]), ..., (a[n-2],a[n-1]).
  // The distance of this zip range is min(n, n-1) = n-1 (via zip_iterator::operator-).
  const auto first_zip = ::cuda::make_zip_iterator(first, first + 1);
  const auto last_zip  = ::cuda::make_zip_iterator(last, last);
  const auto result    = cuda_cub::find_if(policy, first_zip, last_zip, __is_partitioned_fn<Predicate>{predicate});
  // Checking get<1>(result) == last (rather than result == last_zip) correctly
  // handles the n==1 edge case where find_if_n returns first_zip (num_items==0).
  return ::cuda::std::get<1>(result.__iterators()) == last;
}
} // namespace cuda_cub
THRUST_NAMESPACE_END
#endif // _CCCL_CUDA_COMPILATION()
