// SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#include "insert_nested_NVTX_range_guard.h"

#include <cub/device/device_reduce.cuh>

#include <thrust/sequence.h>

#include <cuda/__cmath/uabs.h>
#include <cuda/devices>
#include <cuda/std/__algorithm/max_element.h>
#include <cuda/std/__algorithm/min_element.h>
#include <cuda/std/execution>

#include <cstdint>

#include "catch2_test_device_reduce.cuh"
#include "catch2_test_launch_helper.h"
#include <c2h/catch2_test_helper.h>
#include <c2h/custom_type.h>
#include <c2h/extended_types.h>

DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Reduce, device_reduce);
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Sum, device_sum);
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Min, device_min);
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::ArgMin, device_arg_min);
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Max, device_max);
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::ArgMax, device_arg_max);

_CCCL_SUPPRESS_DEPRECATED_PUSH
_CCCL_SUPPRESS_DEPRECATED_NVRTC_DIAG
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::ArgMin, device_arg_min_old);
DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::ArgMax, device_arg_max_old);
_CCCL_SUPPRESS_DEPRECATED_POP

// %PARAM% TEST_LAUNCH lid 0:1:2
// %PARAM% TEST_TYPES types 0:1:2:3:4

// List of types to test
using custom_t =
  c2h::custom_type_t<c2h::accumulateable_t,
                     c2h::equal_comparable_t,
                     c2h::lexicographical_less_comparable_t,
                     c2h::lexicographical_greater_comparable_t>;

#if TEST_TYPES == 0
using full_type_list = c2h::type_list<type_pair<std::uint8_t>, type_pair<std::int8_t, std::int32_t>>;
#elif TEST_TYPES == 1
using full_type_list = c2h::type_list<type_pair<std::int32_t>, type_pair<std::int64_t>>;
#elif TEST_TYPES == 2
using full_type_list =
  c2h::type_list<type_pair<uchar3>,
                 type_pair<
#  if _CCCL_CTK_AT_LEAST(13, 0)
                   ulonglong4_16a
#  else // _CCCL_CTK_AT_LEAST(13, 0)
                   ulonglong4
#  endif // _CCCL_CTK_AT_LEAST(13, 0)
                   >>;
#elif TEST_TYPES == 3
// clang-format off
using full_type_list = c2h::type_list<
type_pair<custom_t>
#if TEST_HALF_T()
, type_pair<half_t>
#endif // TEST_HALF_T()
#if TEST_BF_T()
, type_pair<bfloat16_t>
#endif // TEST_BF_T()
>;
// clang-format on
#elif TEST_TYPES == 4
// DPX SIMD instructions
using full_type_list = c2h::type_list<type_pair<std::uint16_t>, type_pair<std::int16_t>>;
#endif

/**
 * @brief Input data generation mode
 */
enum class gen_data_t : int
{
  /// Uniform random data generation
  GEN_TYPE_RANDOM,
  /// Constant value as input data
  GEN_TYPE_CONST
};

struct abs_less_t
{
  template <typename T>
  _CCCL_HOST_DEVICE_API auto operator()(const T& a, const T& b) const -> bool
  {
    // need to use `uabs` to avoid integer overflow in case of abs(INT_MIN)
    return cuda::uabs(a) < cuda::uabs(b);
  }
};

C2H_TEST("Device reduce works with all device interfaces", "[reduce][device]", full_type_list)
{
  using params   = params_t<TestType>;
  using item_t   = typename params::item_t;
  using output_t = typename params::output_t;
  using offset_t = int32_t;

  constexpr int max_items    = 5000000;
  constexpr int min_items    = 1;
  constexpr int num_segments = 1;

  // Generate the input sizes to test for
  const int num_items = GENERATE_COPY(
    take(3, random(min_items, max_items)),
    values({
      min_items,
      max_items,
    }));

  // Input data generation to test
  const gen_data_t data_gen_mode = GENERATE_COPY(gen_data_t::GEN_TYPE_RANDOM, gen_data_t::GEN_TYPE_CONST);

  // Generate input data
  c2h::device_vector<item_t> in_items(num_items);
  if (data_gen_mode == gen_data_t::GEN_TYPE_RANDOM)
  {
    c2h::gen(C2H_SEED(2), in_items);
  }
  else
  {
    item_t default_constant{};
    init_default_constant(default_constant);
    thrust::fill(c2h::device_policy, in_items.begin(), in_items.end(), default_constant);
  }
  auto d_in_it = thrust::raw_pointer_cast(in_items.data());

  CAPTURE(c2h::type_name<item_t>(), c2h::type_name<output_t>(), num_items);

#if TEST_TYPES != 4
  SECTION("reduce")
  {
    using op_t = cuda::std::plus<>;

    // Binary reduction operator
    auto reduction_op = unwrap_op(reference_extended_fp(d_in_it), op_t{});

    // Prepare verification data
    using accum_t = cuda::std::__accumulator_t<op_t, item_t, output_t>;
    output_t expected_result =
      static_cast<output_t>(compute_single_problem_reference(in_items, reduction_op, accum_t{}));

    // Run test
    c2h::device_vector<output_t> out_result(num_segments);
    auto d_out_it      = thrust::raw_pointer_cast(out_result.data());
    using init_value_t = cub::detail::it_value_t<decltype(unwrap_it(d_out_it))>;
    device_reduce(unwrap_it(d_in_it), unwrap_it(d_out_it), num_items, reduction_op, init_value_t{});

    // Verify result
    REQUIRE(expected_result == out_result[0]);
  }
#endif // TEST_TYPES != 4

// Skip DeviceReduce::Sum tests for extended floating-point types because of unbounded epsilon due
// to pseudo associativity of the addition operation over floating point numbers
#if TEST_TYPES != 3
  SECTION("sum")
  {
    using op_t    = cuda::std::plus<>;
    using accum_t = cuda::std::__accumulator_t<op_t, item_t, output_t>;

    // Prepare verification data
    output_t expected_result = static_cast<output_t>(compute_single_problem_reference(in_items, op_t{}, accum_t{}));

    // Run test
    c2h::device_vector<output_t> out_result(num_segments);
    auto d_out_it = unwrap_it(thrust::raw_pointer_cast(out_result.data()));
    device_sum(d_in_it, d_out_it, num_items);

    // Verify result
    REQUIRE(expected_result == out_result[0]);
  }
#endif

  SECTION("min")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = *std::min_element(host_items.cbegin(), host_items.cend());

    // Run test
    c2h::device_vector<output_t> out_result(num_segments);
    auto d_out_it = thrust::raw_pointer_cast(out_result.data());
    device_min(unwrap_it(d_in_it), unwrap_it(d_out_it), num_items);

    // Verify result
    REQUIRE(expected_result == out_result[0]);
  }

  SECTION("max")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = *std::max_element(host_items.cbegin(), host_items.cend());

    // Run test
    c2h::device_vector<output_t> out_result(num_segments);
    auto d_out_it = thrust::raw_pointer_cast(out_result.data());
    device_max(unwrap_it(d_in_it), unwrap_it(d_out_it), num_items);

    // Verify result
    REQUIRE(expected_result == out_result[0]);
  }

#if TEST_TYPES != 4
  SECTION("argmax")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = std::max_element(host_items.cbegin(), host_items.cend());

    // Run test
    using result_t = cuda::std::pair<cuda::std::int32_t, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    auto d_result_ptr   = thrust::raw_pointer_cast(out_result.data());
    auto d_index_out    = &d_result_ptr->first;
    auto d_extremum_out = &d_result_ptr->second;
    device_arg_max(unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items);

    // Verify result
    result_t gpu_result   = out_result[0];
    output_t gpu_extremum = static_cast<output_t>(gpu_result.second); // Explicitly rewrap the gpu value
    REQUIRE(expected_result[0] == gpu_extremum);
    REQUIRE((expected_result - host_items.cbegin()) == gpu_result.first);
  }

  SECTION("argmax with user provided memory and environment")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = std::max_element(host_items.cbegin(), host_items.cend());

    // Run test
    using result_t = cuda::std::pair<cuda::std::int32_t, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    auto d_result_ptr   = thrust::raw_pointer_cast(out_result.data());
    auto d_index_out    = &d_result_ptr->first;
    auto d_extremum_out = &d_result_ptr->second;

    size_t expected_allocation_size = 0;
    auto error                      = cub::DeviceReduce::ArgMax(
      static_cast<void*>(nullptr), expected_allocation_size, unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items);
    REQUIRE(error == cudaSuccess);
    REQUIRE(cudaSuccess == cudaPeekAtLastError());
    REQUIRE(cudaSuccess == cudaDeviceSynchronize());

    auto d_temp        = c2h::device_vector<uint8_t>(expected_allocation_size, thrust::no_init);
    void* temp_storage = thrust::raw_pointer_cast(d_temp.data());

    auto test_argmax = [&](const auto& env) {
      size_t num_bytes = 0;
      error            = cub::DeviceReduce::ArgMax(
        static_cast<void*>(nullptr), num_bytes, unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items, env);
      REQUIRE(error == cudaSuccess);
      REQUIRE(cudaSuccess == cudaPeekAtLastError());
      REQUIRE(cudaSuccess == cudaDeviceSynchronize());
      REQUIRE(expected_allocation_size == num_bytes);

      error = cub::DeviceReduce::ArgMax(
        temp_storage, num_bytes, unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items, env);
      REQUIRE(error == cudaSuccess);
      REQUIRE(cudaSuccess == cudaPeekAtLastError());
      REQUIRE(cudaSuccess == cudaDeviceSynchronize());

      // Verify result
      result_t gpu_result   = out_result[0];
      output_t gpu_extremum = static_cast<output_t>(gpu_result.second); // Explicitly rewrap the gpu value
      REQUIRE(expected_result[0] == gpu_extremum);
      REQUIRE((expected_result - host_items.cbegin()) == gpu_result.first);
    };

    int current_device;
    error = cudaGetDevice(&current_device);
    REQUIRE(error == cudaSuccess);

    SECTION("DeviceReduce::ArgMax works with cudaStream_t")
    {
      cuda::stream stream{cuda::devices[current_device]};
      test_argmax(stream.get());
    }

    SECTION("DeviceReduce::ArgMax works with cuda::stream")
    {
      cuda::stream stream{cuda::devices[current_device]};
      test_argmax(stream);
    }

    SECTION("DeviceReduce::ArgMax works with cuda::stream_ref")
    {
      cuda::stream stream{cuda::devices[current_device]};
      cuda::stream_ref stream_ref{stream};
      test_argmax(stream_ref);
    }

    SECTION("DeviceReduce::ArgMax works with cuda::std::execution::env")
    {
      cuda::std::execution::env env{};
      test_argmax(env);
    }

    SECTION("DeviceReduce::ArgMax works with cuda::execution::gpu")
    {
      const auto policy = cuda::execution::gpu;
      test_argmax(policy);
    }

    SECTION("DeviceReduce::ArgMax works with cuda::execution::gpu with stream")
    {
      cuda::stream stream{cuda::devices[current_device]};
      const auto policy = cuda::execution::gpu.with(cuda::get_stream, stream);
      test_argmax(policy);
    }
  }

  SECTION("argmin")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = std::min_element(host_items.cbegin(), host_items.cend());

    // Run test
    using result_t = cuda::std::pair<cuda::std::int32_t, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    auto d_result_ptr   = thrust::raw_pointer_cast(out_result.data());
    auto d_index_out    = &d_result_ptr->first;
    auto d_extremum_out = &d_result_ptr->second;
    device_arg_min(unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items);

    // Verify result
    result_t gpu_result   = out_result[0];
    output_t gpu_extremum = static_cast<output_t>(gpu_result.second); // Explicitly rewrap the gpu value
    REQUIRE(expected_result[0] == gpu_extremum);
    REQUIRE((expected_result - host_items.cbegin()) == gpu_result.first);
  }

  SECTION("argmin with user provided memory and environment")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = std::min_element(host_items.cbegin(), host_items.cend());

    // Run test
    using result_t = cuda::std::pair<cuda::std::int32_t, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    auto d_result_ptr   = thrust::raw_pointer_cast(out_result.data());
    auto d_index_out    = &d_result_ptr->first;
    auto d_extremum_out = &d_result_ptr->second;

    size_t expected_allocation_size = 0;
    auto error                      = cub::DeviceReduce::ArgMin(
      static_cast<void*>(nullptr), expected_allocation_size, unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items);
    REQUIRE(error == cudaSuccess);
    REQUIRE(cudaSuccess == cudaPeekAtLastError());
    REQUIRE(cudaSuccess == cudaDeviceSynchronize());

    auto d_temp        = c2h::device_vector<uint8_t>(expected_allocation_size, thrust::no_init);
    void* temp_storage = thrust::raw_pointer_cast(d_temp.data());

    auto test_argmin = [&](const auto& env) {
      size_t num_bytes = 0;
      error            = cub::DeviceReduce::ArgMin(
        static_cast<void*>(nullptr), num_bytes, unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items, env);
      REQUIRE(error == cudaSuccess);
      REQUIRE(cudaSuccess == cudaPeekAtLastError());
      REQUIRE(cudaSuccess == cudaDeviceSynchronize());
      REQUIRE(expected_allocation_size == num_bytes);

      error = cub::DeviceReduce::ArgMin(
        temp_storage, num_bytes, unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items, env);
      REQUIRE(error == cudaSuccess);
      REQUIRE(cudaSuccess == cudaPeekAtLastError());
      REQUIRE(cudaSuccess == cudaDeviceSynchronize());

      // Verify result
      result_t gpu_result   = out_result[0];
      output_t gpu_extremum = static_cast<output_t>(gpu_result.second); // Explicitly rewrap the gpu value
      REQUIRE(expected_result[0] == gpu_extremum);
      REQUIRE((expected_result - host_items.cbegin()) == gpu_result.first);
    };

    int current_device;
    error = cudaGetDevice(&current_device);
    REQUIRE(error == cudaSuccess);

    SECTION("DeviceReduce::ArgMin works with cudaStream_t")
    {
      cuda::stream stream{cuda::devices[current_device]};
      test_argmin(stream.get());
    }

    SECTION("DeviceReduce::ArgMin works with cuda::stream")
    {
      cuda::stream stream{cuda::devices[current_device]};
      test_argmin(stream);
    }

    SECTION("DeviceReduce::ArgMin works with cuda::stream_ref")
    {
      cuda::stream stream{cuda::devices[current_device]};
      cuda::stream_ref stream_ref{stream};
      test_argmin(stream_ref);
    }

    SECTION("DeviceReduce::ArgMin works with cuda::std::execution::env")
    {
      cuda::std::execution::env env{};
      test_argmin(env);
    }

    SECTION("DeviceReduce::ArgMin works with cuda::execution::gpu")
    {
      const auto policy = cuda::execution::gpu;
      test_argmin(policy);
    }

    SECTION("DeviceReduce::ArgMin works with cuda::execution::gpu with stream")
    {
      cuda::stream stream{cuda::devices[current_device]};
      const auto policy = cuda::execution::gpu.with(cuda::get_stream, stream);
      test_argmin(policy);
    }
  }

  SECTION("argmax deprecated interface")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = std::max_element(host_items.cbegin(), host_items.cend());

    // Run test using the deprecated interface
    using result_t = cub::KeyValuePair<int, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    device_arg_max_old(unwrap_it(d_in_it), thrust::raw_pointer_cast(out_result.data()), num_items);

    // Verify result for the deprecated interface
    result_t gpu_result = out_result[0];
    output_t gpu_value  = static_cast<output_t>(gpu_result.value); // Explicitly rewrap the gpu value
    REQUIRE(expected_result[0] == gpu_value);
    REQUIRE((expected_result - host_items.cbegin()) == gpu_result.key);
  }

  SECTION("argmin deprecated interface")
  {
    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = std::min_element(host_items.cbegin(), host_items.cend());

    // Run test using the deprecated interface
    using result_t = cub::KeyValuePair<int, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    device_arg_min_old(unwrap_it(d_in_it), thrust::raw_pointer_cast(out_result.data()), num_items);

    // Verify result for the deprecated interface
    result_t gpu_result = out_result[0];
    output_t gpu_value  = static_cast<output_t>(gpu_result.value); // Explicitly rewrap the gpu value
    REQUIRE(expected_result[0] == gpu_value);
    REQUIRE((expected_result - host_items.cbegin()) == gpu_result.key);
  }

#  if TEST_TYPES < 2
  SECTION("argmin-abs_less_t")
  {
    abs_less_t compare_op;

    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = cuda::std::min_element(host_items.cbegin(), host_items.cend(), compare_op);

    // Run test
    using result_t = cuda::std::pair<cuda::std::int32_t, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    auto d_result_ptr   = thrust::raw_pointer_cast(out_result.data());
    auto d_index_out    = &d_result_ptr->first;
    auto d_extremum_out = &d_result_ptr->second;
    device_arg_min(unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items, compare_op);

    // Verify result
    result_t gpu_result   = out_result[0];
    output_t gpu_extremum = static_cast<output_t>(gpu_result.second); // Explicitly rewrap the gpu value
    REQUIRE(expected_result[0] == gpu_extremum);
    REQUIRE((expected_result - host_items.cbegin()) == gpu_result.first);
  }

  SECTION("argmax-abs_less_t")
  {
    abs_less_t compare_op;

    // Prepare verification data
    c2h::host_vector<item_t> host_items(in_items);
    auto expected_result = cuda::std::max_element(host_items.cbegin(), host_items.cend(), compare_op);

    // Run test
    using result_t = cuda::std::pair<cuda::std::int32_t, unwrap_value_t<output_t>>;
    c2h::device_vector<result_t> out_result(num_segments);
    auto d_result_ptr   = thrust::raw_pointer_cast(out_result.data());
    auto d_index_out    = &d_result_ptr->first;
    auto d_extremum_out = &d_result_ptr->second;
    device_arg_max(unwrap_it(d_in_it), d_extremum_out, d_index_out, num_items, compare_op);

    // Verify result
    result_t gpu_result   = out_result[0];
    output_t gpu_extremum = static_cast<output_t>(gpu_result.second); // Explicitly rewrap the gpu value
    REQUIRE(expected_result[0] == gpu_extremum);
    REQUIRE((expected_result - host_items.cbegin()) == gpu_result.first);
  }
#  endif
#endif
}

#if TEST_TYPES == 0
// this type stands in for lambda functions, which are also not copy-assignable before C++17
struct non_copy_assignable_plus
{
  non_copy_assignable_plus()                                           = default;
  non_copy_assignable_plus(const non_copy_assignable_plus&)            = default;
  non_copy_assignable_plus& operator=(const non_copy_assignable_plus&) = delete;

  template <typename T>
  _CCCL_HOST_DEVICE_API auto operator()(const T& a, const T& b) const -> T
  {
    return a + b;
  }
};

struct non_copy_assignable_less
{
  non_copy_assignable_less()                                           = default;
  non_copy_assignable_less(const non_copy_assignable_less&)            = default;
  non_copy_assignable_less& operator=(const non_copy_assignable_less&) = delete;

  template <typename T>
  _CCCL_HOST_DEVICE_API auto operator()(const T& a, const T& b) const -> bool
  {
    return a < b;
  }
};

C2H_TEST("Device reduce works with a non copy assignable reduction operator", "[reduce][device]")
{
  using item_t   = int;
  using output_t = int;

  constexpr int num_items = 1000;

  c2h::device_vector<item_t> input(num_items, 42);
  thrust::sequence(input.begin(), input.end(), 1);

  SECTION("reduce")
  {
    c2h::device_vector<output_t> output(1);
    device_reduce(input.data(), output.data(), num_items, non_copy_assignable_plus{}, 0);
    CHECK((num_items * (num_items + 1)) / 2 == output[0]);
  }

  SECTION("argmin")
  {
    c2h::device_vector<output_t> output_extremum(1);
    c2h::device_vector<int> output_index(1);
    device_arg_min(input.data(), output_extremum.data(), output_index.data(), num_items, non_copy_assignable_less{});
    REQUIRE(1 == output_extremum[0]);
    REQUIRE(0 == output_index[0]);
  }
}

struct checking_reduce
{
  static constexpr auto sentinel = 42;

  _CCCL_HOST_DEVICE_API auto operator()(int a, int b) const -> int
  {
    CHECK(a == sentinel);
    CHECK(b == sentinel);
    return sentinel;
  }
};

struct faulting_reduce
{
  _CCCL_HOST_DEVICE_API auto operator()(int, int) const -> int
  {
    CHECK(false);
    return 0;
  }
};

C2H_TEST("Device reduce works without initial value", "[reduce][device]")
{
  constexpr int num_items = 1000;
  c2h::device_vector<int> input(num_items, checking_reduce::sentinel);

  SECTION("for some elements")
  {
    c2h::device_vector<int> output(1);
    device_reduce(input.data(), output.data(), num_items, checking_reduce{}, cub::detail::reduce::no_init);
    CHECK(output[0] == checking_reduce::sentinel);
  }
  SECTION("for no elements")
  {
    device_reduce(input.data(), static_cast<int*>(nullptr), 0, faulting_reduce{}, cub::detail::reduce::no_init);
  }
}
#endif // TEST_TYPES == 0