// 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/internal_functional.h>
#include <thrust/system/detail/generic/count.h>
#include <thrust/transform_reduce.h>

THRUST_NAMESPACE_BEGIN
namespace system::detail::generic
{
template <typename InputType, typename Predicate, typename CountType>
struct count_if_transform
{
  _CCCL_HOST_DEVICE count_if_transform(Predicate _pred)
      : pred(_pred)
  {}

  _CCCL_EXEC_CHECK_DISABLE
  _CCCL_HOST_DEVICE CountType operator()(const InputType& val)
  {
    if (pred(val))
    {
      return 1;
    }
    else
    {
      return 0;
    }
  } // end operator()

  Predicate pred;
}; // end count_if_transform

template <typename DerivedPolicy, typename InputIterator, typename EqualityComparable>
_CCCL_HOST_DEVICE thrust::detail::it_difference_t<InputIterator>
count(thrust::execution_policy<DerivedPolicy>& exec,
      InputIterator first,
      InputIterator last,
      const EqualityComparable& value)
{
  using thrust::placeholders::_1;

  return thrust::count_if(exec, first, last, _1 == value);
} // end count()

template <typename DerivedPolicy, typename InputIterator, typename Predicate>
_CCCL_HOST_DEVICE thrust::detail::it_difference_t<InputIterator>
count_if(thrust::execution_policy<DerivedPolicy>& exec, InputIterator first, InputIterator last, Predicate pred)
{
  using InputType = thrust::detail::it_value_t<InputIterator>;
  using CountType = thrust::detail::it_difference_t<InputIterator>;

  thrust::system::detail::generic::count_if_transform<InputType, Predicate, CountType> unary_op(pred);
  ::cuda::std::plus<CountType> binary_op;
  return thrust::transform_reduce(exec, first, last, unary_op, CountType(0), binary_op);
} // end count_if()
} // namespace system::detail::generic
THRUST_NAMESPACE_END
