// SPDX-FileCopyrightText: Copyright (c) 2008-2013, NVIDIA Corporation. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

#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
#include <thrust/detail/function.h>
#include <thrust/iterator/iterator_traits.h>

#include <cuda/std/__utility/pair.h>

THRUST_NAMESPACE_BEGIN

namespace system::detail::generic::scalar
{
template <typename RandomAccessIterator, typename Size, typename T, typename BinaryPredicate>
_CCCL_HOST_DEVICE RandomAccessIterator
lower_bound_n(RandomAccessIterator first, Size n, const T& val, BinaryPredicate comp)
{
  // wrap comp
  thrust::detail::wrapped_function<BinaryPredicate, bool> wrapped_comp{comp};

  Size start = 0, i;
  while (start < n)
  {
    i = start + (n - start) / 2; // Overflow-safe variant of (a+b)/2
    if (wrapped_comp(first[i], val))
    {
      start = i + 1;
    }
    else
    {
      n = i;
    }
  } // end while

  return first + start;
}

// XXX generalize these upon implementation of scalar::distance & scalar::advance

template <typename RandomAccessIterator, typename T, typename BinaryPredicate>
_CCCL_HOST_DEVICE RandomAccessIterator
lower_bound(RandomAccessIterator first, RandomAccessIterator last, const T& val, BinaryPredicate comp)
{
  thrust::detail::it_difference_t<RandomAccessIterator> n = last - first;
  return lower_bound_n(first, n, val, comp);
}

template <typename RandomAccessIterator, typename Size, typename T, typename BinaryPredicate>
_CCCL_HOST_DEVICE RandomAccessIterator
upper_bound_n(RandomAccessIterator first, Size n, const T& val, BinaryPredicate comp)
{
  // wrap comp
  thrust::detail::wrapped_function<BinaryPredicate, bool> wrapped_comp{comp};

  Size start = 0, i;
  while (start < n)
  {
    i = start + (n - start) / 2; // Overflow-safe variant of (a+b)/2
    if (wrapped_comp(val, first[i]))
    {
      n = i;
    }
    else
    {
      start = i + 1;
    }
  } // end while

  return first + start;
}

template <typename RandomAccessIterator, typename T, typename BinaryPredicate>
_CCCL_HOST_DEVICE RandomAccessIterator
upper_bound(RandomAccessIterator first, RandomAccessIterator last, const T& val, BinaryPredicate comp)
{
  thrust::detail::it_difference_t<RandomAccessIterator> n = last - first;
  return upper_bound_n(first, n, val, comp);
}

template <typename RandomAccessIterator, typename T, typename BinaryPredicate>
_CCCL_HOST_DEVICE ::cuda::std::pair<RandomAccessIterator, RandomAccessIterator>
equal_range(RandomAccessIterator first, RandomAccessIterator last, const T& val, BinaryPredicate comp)
{
  RandomAccessIterator lb = thrust::system::detail::generic::scalar::lower_bound(first, last, val, comp);
  return ::cuda::std::make_pair(lb, thrust::system::detail::generic::scalar::upper_bound(lb, last, val, comp));
}

template <typename RandomAccessIterator, typename T, typename Compare>
_CCCL_HOST_DEVICE bool binary_search(RandomAccessIterator first, RandomAccessIterator last, const T& value, Compare comp)
{
  RandomAccessIterator iter = thrust::system::detail::generic::scalar::lower_bound(first, last, value, comp);

  // wrap comp
  thrust::detail::wrapped_function<Compare, bool> wrapped_comp{comp};

  return iter != last && !wrapped_comp(value, *iter);
}
} // namespace system::detail::generic::scalar

THRUST_NAMESPACE_END

#include <thrust/system/detail/generic/scalar/binary_search.inl>
