// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "insert_nested_NVTX_range_guard.h"

#include <cub/device/device_merge.cuh>

#include <thrust/iterator/zip_iterator.h>
#include <thrust/sort.h>

#include <cuda/iterator>

#include <algorithm>

#include <test_util.h>

#include "catch2_test_launch_helper.h"
#include <c2h/catch2_test_helper.h>

// %PARAM% TEST_LAUNCH lid 0:1:2

DECLARE_LAUNCH_WRAPPER(cub::DeviceMerge::MergePairs, merge_pairs);
DECLARE_LAUNCH_WRAPPER(cub::DeviceMerge::MergeKeys, merge_keys);

using types = c2h::type_list<std::uint8_t, std::int16_t, std::uint32_t, double>;

template <typename Key,
          typename Offset,
          typename CompareOp = cuda::std::less<Key>,
          typename MergeKeys = decltype(::merge_keys)>
void test_keys(Offset size1 = 3623, Offset size2 = 6346, CompareOp compare_op = {}, MergeKeys merge_keys = ::merge_keys)
{
  CAPTURE(c2h::type_name<Key>(), c2h::type_name<Offset>(), size1, size2);

  c2h::device_vector<Key> keys1_d(size1, thrust::default_init);
  c2h::device_vector<Key> keys2_d(size2, thrust::default_init);

  c2h::gen(C2H_SEED(1), keys1_d);
  c2h::gen(C2H_SEED(1), keys2_d);

  thrust::sort(c2h::device_policy, keys1_d.begin(), keys1_d.end(), compare_op);
  thrust::sort(c2h::device_policy, keys2_d.begin(), keys2_d.end(), compare_op);
  // CAPTURE(keys1_d, keys2_d);

  c2h::device_vector<Key> result_d(size1 + size2, thrust::default_init);
  merge_keys(thrust::raw_pointer_cast(keys1_d.data()),
             static_cast<Offset>(keys1_d.size()),
             thrust::raw_pointer_cast(keys2_d.data()),
             static_cast<Offset>(keys2_d.size()),
             thrust::raw_pointer_cast(result_d.data()),
             compare_op);

  c2h::host_vector<Key> keys1_h = keys1_d;
  c2h::host_vector<Key> keys2_h = keys2_d;
  c2h::host_vector<Key> reference_h(size1 + size2, thrust::default_init);
  std::merge(keys1_h.begin(), keys1_h.end(), keys2_h.begin(), keys2_h.end(), reference_h.begin(), compare_op);

  // comparing std::vectors instead compiles in 1m19s, thrust::host_vector 1m23s, thrust::device_vector 1m38
  // let's pick the host_vector, so we don't stress device memory with another (potentially big) allocation
  c2h::host_vector<Key> result_h(result_d); // perform copy outside CHECK() to propagate a potential bad_alloc
  CHECK(reference_h == result_h);
}

C2H_TEST("DeviceMerge::MergeKeys key types", "[merge][device]", types)
{
  using key_t    = c2h::get<0, TestType>;
  using offset_t = int;
  test_keys<key_t, offset_t>();
}

C2H_TEST("DeviceMerge::MergeKeys works for large number of items",
         "[merge][device][skip-cs-racecheck][skip-cs-initcheck][skip-cs-synccheck]")
try
{
  using key_t    = char;
  using offset_t = int64_t;

  // Clamp 64-bit offset type problem sizes to just slightly larger than 2^32 items
  const auto num_items_int_max = static_cast<offset_t>(cuda::std::numeric_limits<std::int32_t>::max());

  // Generate the input sizes to test for
  const offset_t num_items_lhs =
    GENERATE_COPY(values({num_items_int_max + offset_t{1000000}, num_items_int_max - 1, offset_t{3}}));
  const offset_t num_items_rhs =
    GENERATE_COPY(values({num_items_int_max + offset_t{1000000}, num_items_int_max, offset_t{3}}));

  test_keys<key_t, offset_t>(num_items_lhs, num_items_rhs, cuda::std::less<>{});
}
catch (const std::bad_alloc&)
{
  // allocation failure is not a test failure, so we can run tests on smaller GPUs
  SUCCEED("allocation failure is not a test failure");
}

C2H_TEST("DeviceMerge::MergeKeys input sizes", "[merge][device]")
{
  using key_t    = int;
  using offset_t = int;
  // TODO(bgruber): maybe less combinations
  const auto size1 = offset_t{GENERATE(0, 1, 23, 123, 3234)};
  const auto size2 = offset_t{GENERATE(0, 1, 52, 556, 56767)};
  test_keys<key_t>(size1, size2);
}

C2H_TEST("DeviceMerge::MergeKeys almost tile-sized input sizes", "[merge][device]")
{
  using key_t    = int;
  using offset_t = int;

  cuda::compute_capability cc{};
  REQUIRE(cub::detail::ptx_compute_cap(cc) == cudaSuccess);
  const offset_t items_per_tile =
    cub::detail::merge::policy_selector_from_types<key_t*, cub::NullType*, key_t*, cub::NullType*, offset_t>{}(cc)
      .items_per_thread;

  test_keys<key_t>(items_per_tile - 1, 1);
  test_keys<key_t>(items_per_tile, 1);
  test_keys<key_t>(1, items_per_tile - 1);
  test_keys<key_t>(1, items_per_tile);
}

// cannot put those in an anon namespace, or nvcc complains that the kernels have internal linkage
using unordered_t = c2h::custom_type_t<c2h::equal_comparable_t>;
struct order
{
  __host__ __device__ auto operator()(const unordered_t& a, const unordered_t& b) const -> bool
  {
    return a.key < b.key;
  }
};

C2H_TEST("DeviceMerge::MergeKeys no operator<", "[merge][device]")
{
  using key_t    = unordered_t;
  using offset_t = int;
  test_keys<key_t, offset_t, order>();
}

namespace
{
// must use thrust::make_zip_iterator for now
// see https://github.com/NVIDIA/cccl/issues/6400
template <typename... Its>
auto zip(Its... its) -> decltype(thrust::make_zip_iterator(its...))
{
  return thrust::make_zip_iterator(its...);
}

template <typename Value>
struct key_to_value
{
  template <typename Key>
  __host__ __device__ auto operator()(const Key& k) const -> Value
  {
    Value v{};
    convert(k, v, 0);
    return v;
  }

  template <typename Key>
  __host__ __device__ static void convert(const Key& k, Value& v, ...)
  {
    v = static_cast<Value>(k);
  }

  template <template <typename> class... Policies>
  __host__ __device__ static void convert(const c2h::custom_type_t<Policies...>& k, Value& v, int)
  {
    v = static_cast<Value>(k.val);
  }

  template <typename Key, template <typename> class... Policies>
  __host__ __device__ static void convert(const Key& k, c2h::custom_type_t<Policies...>& v, int)
  {
    v     = {};
    v.val = static_cast<decltype(v.val)>(k);
  }
};
} // namespace

template <typename Key,
          typename Value,
          typename Offset,
          typename CompareOp  = cuda::std::less<Key>,
          typename MergePairs = decltype(::merge_pairs)>
void test_pairs(
  Offset size1 = 200, Offset size2 = 625, CompareOp compare_op = {}, MergePairs merge_pairs = ::merge_pairs)
{
  CAPTURE(c2h::type_name<Key>(), c2h::type_name<Value>(), c2h::type_name<Offset>(), size1, size2);

  // we start with random but sorted keys
  c2h::device_vector<Key> keys1_d(size1, thrust::no_init);
  c2h::device_vector<Key> keys2_d(size2, thrust::no_init);
  c2h::gen(C2H_SEED(1), keys1_d);
  c2h::gen(C2H_SEED(1), keys2_d);
  thrust::sort(c2h::device_policy, keys1_d.begin(), keys1_d.end(), compare_op);
  thrust::sort(c2h::device_policy, keys2_d.begin(), keys2_d.end(), compare_op);

  // the values must be functionally dependent on the keys (equal key => equal value), since merge is unstable
  c2h::device_vector<Value> values1_d(size1, thrust::no_init);
  c2h::device_vector<Value> values2_d(size2, thrust::no_init);
  thrust::transform(c2h::device_policy, keys1_d.begin(), keys1_d.end(), values1_d.begin(), key_to_value<Value>{});
  thrust::transform(c2h::device_policy, keys2_d.begin(), keys2_d.end(), values2_d.begin(), key_to_value<Value>{});
  //  CAPTURE(keys1_d, keys2_d, values1_d, values2_d);

  // compute CUB result
  c2h::device_vector<Key> result_keys_d(size1 + size2, thrust::no_init);
  c2h::device_vector<Value> result_values_d(size1 + size2, thrust::no_init);
  merge_pairs(
    thrust::raw_pointer_cast(keys1_d.data()),
    thrust::raw_pointer_cast(values1_d.data()),
    static_cast<Offset>(keys1_d.size()),
    thrust::raw_pointer_cast(keys2_d.data()),
    thrust::raw_pointer_cast(values2_d.data()),
    static_cast<Offset>(keys2_d.size()),
    thrust::raw_pointer_cast(result_keys_d.data()),
    thrust::raw_pointer_cast(result_values_d.data()),
    compare_op);

  // compute reference result
  c2h::host_vector<Key> reference_keys_h(size1 + size2, thrust::no_init);
  c2h::host_vector<Value> reference_values_h(size1 + size2, thrust::no_init);
  {
    c2h::host_vector<Key> keys1_h     = keys1_d;
    c2h::host_vector<Value> values1_h = values1_d;
    c2h::host_vector<Key> keys2_h     = keys2_d;
    c2h::host_vector<Value> values2_h = values2_d;
    using value_t                     = typename decltype(zip(keys1_h.begin(), values1_h.begin()))::value_type;
    std::merge(zip(keys1_h.begin(), values1_h.begin()),
               zip(keys1_h.end(), values1_h.end()),
               zip(keys2_h.begin(), values2_h.begin()),
               zip(keys2_h.end(), values2_h.end()),
               zip(reference_keys_h.begin(), reference_values_h.begin()),
               [&](const value_t& a, const value_t& b) {
                 return compare_op(cuda::std::get<0>(a), cuda::std::get<0>(b));
               });
  }

  // FIXME(bgruber): comparing std::vectors (slower than thrust vectors) but compiles a lot faster
  CHECK((detail::to_vec(reference_keys_h) == detail::to_vec(c2h::host_vector<Key>(result_keys_d))));
  CHECK((detail::to_vec(reference_values_h) == detail::to_vec(c2h::host_vector<Value>(result_values_d))));
}

C2H_TEST("DeviceMerge::MergePairs key types", "[merge][device]", types)
{
  using key_t    = c2h::get<0, TestType>;
  using value_t  = int;
  using offset_t = int;
  test_pairs<key_t, value_t, offset_t>();
}

// TODO(bgruber): fine tune the type sizes again to hit the fallback and the vsmem policies
// C2H_TEST("DeviceMerge::MergePairs large key types", "[merge][device]", large_types)
// {
//   using key_t    = c2h::get<0, TestType>;
//   using value_t  = int;
//   using offset_t = int;
//   test_pairs<key_t, value_t, offset_t>();
// }

C2H_TEST("DeviceMerge::MergePairs value types", "[merge][device]", types)
{
  using key_t    = int;
  using value_t  = c2h::get<0, TestType>;
  using offset_t = int;
  test_pairs<key_t, value_t, offset_t>();
}

C2H_TEST("DeviceMerge::MergePairs input sizes", "[merge][device]")
{
  using key_t      = int;
  using value_t    = int;
  using offset_t   = int;
  const auto size1 = offset_t{GENERATE(0, 1, 23, 123, 3234234)};
  const auto size2 = offset_t{GENERATE(0, 1, 52, 556, 56767)};
  test_pairs<key_t, value_t>(size1, size2);
}

// this test exceeds 4GiB of memory and the range of 32-bit integers
C2H_TEST("DeviceMerge::MergePairs really large input",
         "[merge][device][skip-cs-racecheck][skip-cs-initcheck][skip-cs-synccheck]")
try
{
  using key_t     = char;
  using value_t   = char;
  const auto size = std::int64_t{1} << GENERATE(30, 31, 32, 33);
  test_pairs<key_t, value_t>(size, size, cuda::std::less<>{});
}
catch (const std::bad_alloc&)
{
  // allocation failure is not a test failure, so we can run tests on smaller GPUs
  SUCCEED("allocation failure is not a test failure");
}

C2H_TEST("DeviceMerge::MergePairs iterators", "[merge][device]")
{
  using key_t             = int;
  using value_t           = int;
  using offset_t          = int;
  const offset_t size1    = 363;
  const offset_t size2    = 634;
  const auto values_start = 123456789;

  const auto larger_size  = std::max(size1, size2);
  const auto smaller_size = std::min(size1, size2);

  auto test = [&](auto key1_it, auto value1_it, auto key2_it, auto value2_it) {
    // compute CUB result
    c2h::device_vector<key_t> result_keys_d(size1 + size2);
    c2h::device_vector<value_t> result_values_d(size1 + size2);
    merge_pairs(
      key1_it,
      value1_it,
      size1,
      key2_it,
      value2_it,
      size2,
      result_keys_d.begin(),
      result_values_d.begin(),
      cuda::std::less<key_t>{});

    // check result
    c2h::host_vector<key_t> result_keys_h     = result_keys_d;
    c2h::host_vector<value_t> result_values_h = result_values_d;

    for (offset_t i = 0; i < static_cast<offset_t>(result_keys_h.size()); i++)
    {
      CAPTURE(i);
      if (i < 2 * smaller_size)
      {
        CHECK(result_keys_h[i + 0] == i / 2);
        CHECK(result_values_h[i + 0] == values_start + i / 2);
      }
      else
      {
        CHECK(result_keys_h[i] == i - smaller_size);
        CHECK(result_values_h[i] == values_start + i - smaller_size);
      }
    }
  };

  auto key_it   = cuda::counting_iterator<key_t>{};
  auto value_it = cuda::counting_iterator<key_t>{values_start};

  c2h::device_vector<key_t> keys_vec(larger_size);
  thrust::sequence(keys_vec.begin(), keys_vec.end());
  c2h::device_vector<key_t> values_vec(larger_size);
  thrust::sequence(values_vec.begin(), values_vec.end(), values_start);

  SECTION("cit/cit/cit/cit")
  {
    test(key_it, value_it, key_it, value_it);
  }
  // key arrays have mixed types
  SECTION("vec/cit/cit/cit")
  {
    test(keys_vec.begin(), value_it, key_it, value_it);
  }
  // value arrays have mixed types
  SECTION("cit/vec/cit/cit")
  {
    test(key_it, values_vec.begin(), key_it, value_it);
  }
  // key and value arrays have mixed types
  SECTION("cit/vec/vec/cit")
  {
    test(key_it, values_vec.begin(), keys_vec.begin(), value_it);
  }
  // values have different iterator and keys
  SECTION("cit/vec/cit/vec")
  {
    test(key_it, values_vec.begin(), key_it, values_vec.begin());
  }
}