/******************************************************************************
 * Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
 * Modifications Copyright© 2019-2025 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 ******************************************************************************/
#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 THRUST_DEVICE_COMPILER == THRUST_DEVICE_COMPILER_HIP
#  include <thrust/system/hip/config.h>

#  include <thrust/detail/alignment.h>
#  include <thrust/detail/integer_math.h>
#  include <thrust/detail/temporary_array.h>
#  include <thrust/detail/trivial_sequence.h>
#  include <thrust/distance.h>
#  include <thrust/extrema.h>
#  include <thrust/sequence.h>
#  include <thrust/sort.h>
#  include <thrust/system/hip/detail/execution_policy.h>
#  include <thrust/system/hip/detail/par_to_seq.h>
#  include <thrust/system/hip/detail/util.h>
#  include <thrust/type_traits/is_contiguous_iterator.h>

#  include <cstdint>
#  include <functional>
#  include <type_traits>

// rocPRIM includes
#  include <rocprim/rocprim.hpp>

THRUST_NAMESPACE_BEGIN
namespace hip_rocprim
{

namespace __merge_sort
{

template <class SORT_ITEMS>
struct dispatch;

// sort keys
template <>
struct dispatch<::thrust::detail::false_type>
{
  template <class KeysIt, class ItemsIt, class Size, class CompareOp>
  static hipError_t THRUST_HIP_RUNTIME_FUNCTION doit(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    KeysIt keys,
    ItemsIt /*items*/,
    Size count,
    CompareOp compare_op,
    hipStream_t stream,
    bool debug_sync)
  {
    return rocprim::merge_sort(
      d_temp_storage, temp_storage_bytes, keys, keys, static_cast<unsigned int>(count), compare_op, stream, debug_sync);
  }
};

// sort pairs
template <>
struct dispatch<::thrust::detail::true_type>
{
  template <class KeysIt, class ItemsIt, class Size, class CompareOp>
  static hipError_t THRUST_HIP_RUNTIME_FUNCTION doit(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    KeysIt keys,
    ItemsIt items,
    Size count,
    CompareOp compare_op,
    hipStream_t stream,
    bool debug_sync)
  {
    return rocprim::merge_sort(
      d_temp_storage,
      temp_storage_bytes,
      keys,
      keys,
      items,
      items,
      static_cast<unsigned int>(count),
      compare_op,
      stream,
      debug_sync);
  }
};

template <typename SORT_ITEMS, typename Derived, typename KeysIt, typename ItemsIt, typename CompareOp>
THRUST_HIP_RUNTIME_FUNCTION void merge_sort(
  execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ItemsIt items_first, CompareOp compare_op)
{
  using size_type = typename iterator_traits<KeysIt>::difference_type;

  size_type count = static_cast<size_type>(thrust::distance(keys_first, keys_last));

  if (count == 0)
  {
    return;
  }

  size_t storage_size = 0;
  hipStream_t stream  = hip_rocprim::stream(policy);
  bool debug_sync     = THRUST_HIP_DEBUG_SYNC_FLAG;

  hipError_t status;
  status =
    dispatch<SORT_ITEMS>::doit(nullptr, storage_size, keys_first, items_first, count, compare_op, stream, debug_sync);
  hip_rocprim::throw_on_error(status, "merge_sort: failed on 1st step");

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

  status =
    dispatch<SORT_ITEMS>::doit(ptr, storage_size, keys_first, items_first, count, compare_op, stream, debug_sync);
  hip_rocprim::throw_on_error(status, "merge_sort: failed on 2nd step");

  status = hip_rocprim::synchronize_optional(policy);
  hip_rocprim::throw_on_error(status, "merge_sort: failed to synchronize");
}
} // namespace __merge_sort

namespace __radix_sort
{

template <class SORT_ITEMS, class Comparator>
struct dispatch;

// sort keys in ascending order
template <class KeyOrVoid>
struct dispatch<thrust::detail::false_type, thrust::less<KeyOrVoid>>
{
  template <class KeysIt, class ItemsIt, class Size>
  static hipError_t THRUST_HIP_RUNTIME_FUNCTION doit(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    KeysIt keys,
    ItemsIt /*items*/,
    Size count,
    hipStream_t stream,
    bool debug_sync)
  {
    using KeysType = typename std::iterator_traits<KeysIt>::value_type;
    return rocprim::radix_sort_keys(
      d_temp_storage,
      temp_storage_bytes,
      keys,
      keys,
      static_cast<unsigned int>(count),
      0,
      sizeof(KeysType) * 8,
      stream,
      debug_sync);
  }
}; // struct dispatch -- sort keys in ascending order;

// sort keys in descending order
template <class KeyOrVoid>
struct dispatch<thrust::detail::false_type, thrust::greater<KeyOrVoid>>
{
  template <class KeysIt, class ItemsIt, class Size>
  static hipError_t THRUST_HIP_RUNTIME_FUNCTION doit(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    KeysIt keys,
    ItemsIt /*items*/,
    Size count,
    hipStream_t stream,
    bool debug_sync)
  {
    using KeysType = typename std::iterator_traits<KeysIt>::value_type;
    return rocprim::radix_sort_keys_desc(
      d_temp_storage,
      temp_storage_bytes,
      keys,
      keys,
      static_cast<unsigned int>(count),
      0,
      sizeof(KeysType) * 8,
      stream,
      debug_sync);
  }
}; // struct dispatch -- sort keys in descending order;

// sort pairs in ascending order
template <class KeyOrVoid>
struct dispatch<thrust::detail::true_type, thrust::less<KeyOrVoid>>
{
  template <class KeysIt, class ItemsIt, class Size>
  static hipError_t THRUST_HIP_RUNTIME_FUNCTION doit(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    KeysIt keys,
    ItemsIt items,
    Size count,
    hipStream_t stream,
    bool debug_sync)
  {
    using KeysType = typename std::iterator_traits<KeysIt>::value_type;
    return rocprim::radix_sort_pairs(
      d_temp_storage,
      temp_storage_bytes,
      keys,
      keys,
      items,
      items,
      static_cast<unsigned int>(count),
      0,
      sizeof(KeysType) * 8,
      stream,
      debug_sync);
  }
}; // struct dispatch -- sort pairs in ascending order;

// sort pairs in descending order
template <class KeyOrVoid>
struct dispatch<thrust::detail::true_type, thrust::greater<KeyOrVoid>>
{
  template <class KeysIt, class ItemsIt, class Size>
  static hipError_t THRUST_HIP_RUNTIME_FUNCTION doit(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    KeysIt keys,
    ItemsIt items,
    Size count,
    hipStream_t stream,
    bool debug_sync)
  {
    using KeysType = typename std::iterator_traits<KeysIt>::value_type;
    return rocprim::radix_sort_pairs_desc(
      d_temp_storage,
      temp_storage_bytes,
      keys,
      keys,
      items,
      items,
      static_cast<unsigned int>(count),
      0,
      sizeof(KeysType) * 8,
      stream,
      debug_sync);
  }
}; // struct dispatch -- sort pairs in descending order;

template <class SORT_ITEMS, class KeyOrVoid>
struct dispatch<SORT_ITEMS, ::std::less<KeyOrVoid>> : dispatch<SORT_ITEMS, thrust::less<KeyOrVoid>>
{};

template <class SORT_ITEMS, class KeyOrVoid>
struct dispatch<SORT_ITEMS, ::std::greater<KeyOrVoid>> : dispatch<SORT_ITEMS, thrust::greater<KeyOrVoid>>
{};

template <typename SORT_ITEMS, typename Derived, typename KeysIt, typename ItemsIt, typename CompareOp>
THRUST_HIP_RUNTIME_FUNCTION void
radix_sort(execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ItemsIt items_first, CompareOp)
{
  using size_type = typename iterator_traits<KeysIt>::difference_type;

  const size_type count = static_cast<size_type>(thrust::distance(keys_first, keys_last));

  if (count == 0)
  {
    return;
  }

  size_t storage_size = 0;
  hipStream_t stream  = hip_rocprim::stream(policy);
  bool debug_sync     = THRUST_HIP_DEBUG_SYNC_FLAG;

  hipError_t status;

  status =
    dispatch<SORT_ITEMS, CompareOp>::doit(nullptr, storage_size, keys_first, items_first, count, stream, debug_sync);
  hip_rocprim::throw_on_error(status, "radix_sort: failed on 1st step");

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

  status = dispatch<SORT_ITEMS, CompareOp>::doit(ptr, storage_size, keys_first, items_first, count, stream, debug_sync);
  hip_rocprim::throw_on_error(status, "radix_sort: failed on 2nd step");
  hip_rocprim::throw_on_error(hip_rocprim::synchronize_optional(policy), "radix sort: failed to synchronize");
}
} // namespace __radix_sort

//---------------------------------------------------------------------
// Smart sort picks at compile-time whether to dispatch radix or merge sort
//---------------------------------------------------------------------

namespace __smart_sort
{

// TODO(bgruber): we can drop thrust::less etc. when they truly alias to the ::std ones
template <class Key, class CompareOp>
using can_use_primitive_sort = ::std::integral_constant<
  bool,
  ::std::is_arithmetic<Key>::value
    && (::std::is_same<CompareOp, thrust::less<Key>>::value || ::std::is_same<CompareOp, ::std::less<Key>>::value
        || ::std::is_same<CompareOp, thrust::less<void>>::value || ::std::is_same<CompareOp, ::std::less<void>>::value
        || ::std::is_same<CompareOp, thrust::greater<Key>>::value || ::std::is_same<CompareOp, ::std::greater<Key>>::value
        || ::std::is_same<CompareOp, thrust::greater<void>>::value
        || ::std::is_same<CompareOp, ::std::greater<void>>::value)>;

template <class SORT_ITEMS,
          class Derived,
          class Policy,
          class KeysIt,
          class ItemsIt,
          class CompareOp,
          ::std::enable_if_t<!can_use_primitive_sort<typename iterator_value<KeysIt>::type, CompareOp>::value, int> = 0>
THRUST_HIP_RUNTIME_FUNCTION void
smart_sort(Policy& policy, KeysIt keys_first, KeysIt keys_last, ItemsIt items_first, CompareOp compare_op)
{
  __merge_sort::merge_sort<SORT_ITEMS>(policy, keys_first, keys_last, items_first, compare_op);
}

template <class SORT_ITEMS,
          class /*Derived*/,
          class Policy,
          class KeysIt,
          class ItemsIt,
          class CompareOp,
          ::std::enable_if_t<can_use_primitive_sort<typename iterator_value<KeysIt>::type, CompareOp>::value, int> = 0>
THRUST_HIP_RUNTIME_FUNCTION void smart_sort(
  execution_policy<Policy>& policy, KeysIt keys_first, KeysIt keys_last, ItemsIt items_first, CompareOp compare_op)
{
  __radix_sort::radix_sort<SORT_ITEMS>(policy, keys_first, keys_last, items_first, compare_op);
}
} // namespace __smart_sort

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

THRUST_EXEC_CHECK_DISABLE
template <class Derived, class ItemsIt, class CompareOp>
void THRUST_HOST_DEVICE stable_sort(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last, CompareOp compare_op)
{
  // struct workaround is required for HIP-clang
  struct workaround
  {
    THRUST_HOST static void par(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last, CompareOp compare_op)
    {
      using item_t  = thrust::iterator_value_t<ItemsIt>;
      item_t* null_ = nullptr;
      __smart_sort::smart_sort<thrust::detail::false_type, thrust::detail::false_type>(
        policy, first, last, null_, compare_op);
    }
#  if !__THRUST_HAS_HIPRT__
    THRUST_DEVICE static void seq(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last, CompareOp compare_op)
    {
      thrust::stable_sort(cvt_to_seq(derived_cast(policy)), first, last, compare_op);
    }
#  endif
  };
#  if __THRUST_HAS_HIPRT__
  workaround::par(policy, first, last, compare_op);
#  else
  workaround::seq(policy, first, last, compare_op);
#  endif
}

THRUST_EXEC_CHECK_DISABLE
template <class Derived, class ItemsIt, class CompareOp>
void THRUST_HOST_DEVICE sort(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last, CompareOp compare_op)
{
  hip_rocprim::stable_sort(policy, first, last, compare_op);
}

THRUST_EXEC_CHECK_DISABLE
template <class Derived, class KeysIt, class ValuesIt, class CompareOp>
void THRUST_HOST_DEVICE stable_sort_by_key(
  execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ValuesIt values, CompareOp compare_op)
{
  // struct workaround is required for HIP-clang
  struct workaround
  {
    THRUST_HOST static void
    par(execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ValuesIt values, CompareOp compare_op)
    {
      __smart_sort::smart_sort<thrust::detail::true_type, thrust::detail::false_type>(
        policy, keys_first, keys_last, values, compare_op);
    }
#  if !__THRUST_HAS_HIPRT__
    THRUST_DEVICE static void
    seq(execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ValuesIt values, CompareOp compare_op)
    {
      thrust::stable_sort_by_key(cvt_to_seq(derived_cast(policy)), keys_first, keys_last, values, compare_op);
    }
#  endif
  };

#  if __THRUST_HAS_HIPRT__
  workaround::par(policy, keys_first, keys_last, values, compare_op);
#  else
  workaround::seq(policy, keys_first, keys_last, values, compare_op);
#  endif
}

THRUST_EXEC_CHECK_DISABLE
template <class Derived, class KeysIt, class ValuesIt, class CompareOp>
void THRUST_HOST_DEVICE sort_by_key(
  execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ValuesIt values, CompareOp compare_op)
{
  hip_rocprim::stable_sort_by_key(policy, keys_first, keys_last, values, compare_op);
}

// API with default comparator

template <class Derived, class ItemsIt>
void THRUST_HOST_DEVICE sort(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last)
{
  using item_type = typename thrust::iterator_value<ItemsIt>::type;
  hip_rocprim::sort(policy, first, last, less<item_type>());
}

template <class Derived, class ItemsIt>
void THRUST_HOST_DEVICE stable_sort(execution_policy<Derived>& policy, ItemsIt first, ItemsIt last)
{
  using item_type = typename thrust::iterator_value<ItemsIt>::type;
  hip_rocprim::stable_sort(policy, first, last, less<item_type>());
}

template <class Derived, class KeysIt, class ValuesIt>
void THRUST_HOST_DEVICE
sort_by_key(execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ValuesIt values)
{
  using key_type = typename thrust::iterator_value<KeysIt>::type;
  hip_rocprim::sort_by_key(policy, keys_first, keys_last, values, less<key_type>());
}

template <class Derived, class KeysIt, class ValuesIt>
void THRUST_HOST_DEVICE
stable_sort_by_key(execution_policy<Derived>& policy, KeysIt keys_first, KeysIt keys_last, ValuesIt values)
{
  using key_type = typename thrust::iterator_value<KeysIt>::type;
  hip_rocprim::stable_sort_by_key(policy, keys_first, keys_last, values, less<key_type>());
}

} // namespace hip_rocprim
THRUST_NAMESPACE_END
#endif
