#include <cstdint>
#include <cstdlib>
#include <numeric>
#include <optional> // std::optional
#include <string>

#include <cuda_runtime.h>

#include "algorithm_execution.h"
#include "build_result_caching.h"
#include "test_util.h"
#include <cccl/c/transform.h>
#include <cccl/c/types.h>

using BuildResultT = cccl_device_transform_build_result_t;

struct transform_cleanup
{
  CUresult operator()(BuildResultT* build_data) const noexcept
  {
    return cccl_device_transform_cleanup(build_data);
  }
};

using transform_deleter       = BuildResultDeleter<BuildResultT, transform_cleanup>;
using transform_build_cache_t = build_cache_t<std::string, result_wrapper_t<BuildResultT, transform_deleter>>;

template <typename Tag>
auto& get_cache()
{
  return fixture<transform_build_cache_t, Tag>::get_or_create().get_value();
}

struct transform_build
{
  using IterT = cccl_iterator_t;

  template <typename... Ts>
  CUresult operator()(BuildResultT* build_ptr, IterT input, IterT output, uint64_t, Ts... rest) const noexcept
  {
    return cccl_device_unary_transform_build(build_ptr, input, output, rest...);
  }

  template <typename... Ts>
  CUresult
  operator()(BuildResultT* build_ptr, IterT input1, IterT input2, IterT output, uint64_t, Ts... rest) const noexcept
  {
    return cccl_device_binary_transform_build(build_ptr, input1, input2, output, rest...);
  }
};

struct unary_transform_run
{
  template <typename... Ts>
  CUresult operator()(BuildResultT build, void* scratch, size_t* scratch_size, Ts... args) const noexcept
  {
    *scratch_size = 1;
    return (scratch) ? cccl_device_unary_transform(build, args...) : CUDA_SUCCESS;
  }
};

struct binary_transform_run
{
  template <typename... Ts>
  CUresult operator()(BuildResultT build, void* scratch, size_t* scratch_size, Ts... args) const noexcept
  {
    *scratch_size = 1;
    return (scratch) ? cccl_device_binary_transform(build, args...) : CUDA_SUCCESS;
  }
};

template <typename BuildCache = transform_build_cache_t, typename KeyT = std::string>
void unary_transform(
  cccl_iterator_t input,
  cccl_iterator_t output,
  uint64_t num_items,
  cccl_op_t op,
  std::optional<BuildCache>& cache,
  const std::optional<KeyT>& lookup_key)
{
  AlgorithmExecute<BuildResultT, transform_build, transform_cleanup, unary_transform_run, BuildCache, KeyT>(
    cache, lookup_key, input, output, num_items, op);
}

template <typename BuildCache = transform_build_cache_t, typename KeyT = std::string>
void binary_transform(
  cccl_iterator_t input1,
  cccl_iterator_t input2,
  cccl_iterator_t output,
  uint64_t num_items,
  cccl_op_t op,
  std::optional<BuildCache>& cache,
  const std::optional<KeyT>& lookup_key)
{
  AlgorithmExecute<BuildResultT, transform_build, transform_cleanup, binary_transform_run, BuildCache, KeyT>(
    cache, lookup_key, input1, input2, output, num_items, op);
}

C2H_TEST("Transform generates UBLKCP on SM90", "[transform][ublkcp]")
{
  constexpr int device_id = 0;
  const auto& build_info  = BuildInformation<device_id>::init();

  // Only test for ublkcp when it is actually possible to get it.
  if (build_info.get_cc_major() < 9)
  {
    return;
  }

  cccl_device_transform_build_result_t build{};
  operation_t op = make_operation("op", get_unary_op(get_type_info<int>().type));
  REQUIRE(
    CUDA_SUCCESS
    == cccl_device_unary_transform_build(
      &build,
      pointer_t<int>(0),
      pointer_t<int>(0),
      op,
      build_info.get_cc_major(),
      build_info.get_cc_minor(),
      build_info.get_cub_path(),
      build_info.get_thrust_path(),
      build_info.get_libcudacxx_path(),
      build_info.get_ctk_path()));

  std::string sass = inspect_sass(build.payload, build.payload_size);
  CHECK(sass.find("UBLKCP") != std::string::npos);

  op = make_operation("op", get_reduce_op(get_type_info<int>().type));
  REQUIRE(
    CUDA_SUCCESS
    == cccl_device_binary_transform_build(
      &build,
      pointer_t<int>(0),
      pointer_t<int>(0),
      pointer_t<int>(0),
      op,
      build_info.get_cc_major(),
      build_info.get_cc_minor(),
      build_info.get_cub_path(),
      build_info.get_thrust_path(),
      build_info.get_libcudacxx_path(),
      build_info.get_ctk_path()));

  sass = inspect_sass(build.payload, build.payload_size);
  CHECK(sass.find("UBLKCP") != std::string::npos);
}

using integral_types = c2h::type_list<int32_t, uint32_t, int64_t, uint64_t>;
struct Transform_IntegralTypes_Fixture_Tag;
C2H_TEST("Transform works with integral types", "[transform]", integral_types)
{
  using T = c2h::get<0, TestType>;

  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  operation_t op              = make_operation("op", get_unary_op(get_type_info<T>().type));
  const std::vector<T> input  = generate<T>(num_items);
  const std::vector<T> output(num_items, 0);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr(output);

  auto& build_cache    = get_cache<Transform_IntegralTypes_Fixture_Tag>();
  const auto& test_key = make_key<T>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<T> expected(num_items, 0);
  std::transform(input.begin(), input.end(), expected.begin(), [](const T& x) {
    return 2 * x;
  });

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<T>(output_ptr));
  }
}

struct Transform_MisalignedInput_IntegerTypes_Fixture_Tag;
C2H_TEST("Transform works with misaligned input with integral types", "[transform]", integral_types)
{
  using T = c2h::get<0, TestType>;

  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  operation_t op              = make_operation("op", get_unary_op(get_type_info<T>().type));
  const std::vector<T> input  = generate<T>(num_items + 1);
  const std::vector<T> output(num_items, 0);
  pointer_t<T> input_ptr_aligned(input);
  pointer_t<T> input_ptr = input;
  input_ptr.ptr += 1; // misalign by 1 from the guaranteed alignment of cudaMalloc, to maybe trip vectorized path
  input_ptr.size -= 1;
  pointer_t<T> output_ptr(output);

  auto& build_cache    = get_cache<Transform_MisalignedInput_IntegerTypes_Fixture_Tag>();
  const auto& test_key = make_key<T>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);
  input_ptr.ptr = nullptr; // avoid freeing the memory through this pointer

  std::vector<T> expected(num_items, 0);
  std::transform(input.begin() + 1, input.end(), expected.begin(), [](const T& x) {
    return 2 * x;
  });

  REQUIRE(expected == std::vector<T>(output_ptr));
}

struct Transform_MisalignedOutput_IntegerTypes_Fixture_Tag;
C2H_TEST("Transform works with misaligned output with integral types", "[transform]", integral_types)
{
  using T = c2h::get<0, TestType>;

  const std::size_t num_items = GENERATE(1, 42, take(4, random(1 << 12, 1 << 16)));
  operation_t op              = make_operation("op", get_unary_op(get_type_info<T>().type));
  const std::vector<T> input  = generate<T>(num_items);
  const std::vector<T> output(num_items + 1, 0);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr_aligned(output);
  pointer_t<T> output_ptr = output;
  output_ptr.ptr += 1; // misalign by 1 from the guaranteed alignment of cudaMalloc, to maybe trip vectorized path
  output_ptr.size -= 1;

  auto& build_cache    = get_cache<Transform_MisalignedOutput_IntegerTypes_Fixture_Tag>();
  const auto& test_key = make_key<T>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<T> expected(num_items, 0);
  std::transform(input.begin(), input.end(), expected.begin(), [](const T& x) {
    return 2 * x;
  });

  REQUIRE(expected == std::vector<T>(output_ptr));

  output_ptr.ptr = nullptr; // avoid freeing the memory through this pointer
}

struct Transform_IntegralTypes_WellKnown_Fixture_Tag;
C2H_TEST("Transform works with integral types with well-known operations", "[transform][well_known]", integral_types)
{
  using T = c2h::get<0, TestType>;

  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  cccl_op_t op                = make_well_known_unary_operation();
  const std::vector<T> input  = generate<T>(num_items);
  const std::vector<T> output(num_items, 0);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr(output);

  auto& build_cache    = get_cache<Transform_IntegralTypes_WellKnown_Fixture_Tag>();
  const auto& test_key = make_key<T>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<T> expected(num_items, 0);
  _CCCL_DIAG_PUSH
  _CCCL_DIAG_SUPPRESS_MSVC(4146) // unary minus on unsigned type
  std::transform(input.begin(), input.end(), expected.begin(), [](const T& x) {
    return -x;
  });
  _CCCL_DIAG_POP
  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<T>(output_ptr));
  }
}

C2H_TEST("Transform works with logical and bitwise well-known operations", "[transform][well_known]")
{
  std::optional<transform_build_cache_t> no_cache = std::nullopt;
  const std::optional<std::string> no_key         = std::nullopt;

  {
    const std::vector<uint32_t> input{0U, 1U, 0xaaaaaaaaU, 0xffffffffU};
    pointer_t<uint32_t> input_ptr(input);
    pointer_t<uint32_t> output_ptr(input.size());

    cccl_op_t bit_not_op = make_well_known_unary_operation();
    bit_not_op.type      = cccl_op_kind_t::CCCL_BIT_NOT;
    unary_transform(input_ptr, output_ptr, input.size(), bit_not_op, no_cache, no_key);

    REQUIRE(std::vector<uint32_t>(output_ptr) == std::vector<uint32_t>{0xffffffffU, 0xfffffffeU, 0x55555555U, 0U});
  }

  const std::vector<uint8_t> lhs{1, 1, 0, 0};
  const std::vector<uint8_t> rhs{1, 0, 1, 0};
  pointer_t<uint8_t> lhs_ptr(lhs);
  pointer_t<uint8_t> rhs_ptr(rhs);

  const auto check_logical_op = [&](cccl_op_kind_t kind, const std::vector<uint8_t>& expected) {
    pointer_t<uint8_t> output_ptr(lhs.size());
    cccl_op_t op = make_well_known_binary_operation();
    op.type      = kind;

    binary_transform(
      make_boolean_iterator(lhs_ptr),
      make_boolean_iterator(rhs_ptr),
      make_boolean_iterator(output_ptr),
      lhs.size(),
      op,
      no_cache,
      no_key);

    REQUIRE(std::vector<uint8_t>(output_ptr) == expected);
  };

  check_logical_op(cccl_op_kind_t::CCCL_LOGICAL_AND, {1, 0, 0, 0});
  check_logical_op(cccl_op_kind_t::CCCL_LOGICAL_OR, {1, 1, 1, 0});
}

struct pair
{
  short a;
  size_t b;

  bool operator==(const pair& other) const
  {
    return a == other.a && b == other.b;
  }
};

struct custom_int
{
  int value;
};

C2H_TEST("Transform works with C++ source for custom types with a well-known unary operation",
         "[transform][well_known][cpp_source]")
{
  const std::string source = R"(
struct custom_int { int value; };
extern "C" __device__ void logical_not_custom_int(void* input_ptr, void* output_ptr) {
  const custom_int* input = static_cast<const custom_int*>(input_ptr);
  bool* output = static_cast<bool*>(output_ptr);
  *output = !input->value;
}
)";
  const std::vector<custom_int> input{{0}, {1}, {-2}, {42}};
  pointer_t<custom_int> input_ptr(input);
  std::optional<transform_build_cache_t> no_cache = std::nullopt;
  const std::optional<std::string> no_key         = std::nullopt;

  operation_t op_state = make_cpp_operation("logical_not_custom_int", source);
  cccl_op_t op         = op_state;
  op.type              = cccl_op_kind_t::CCCL_LOGICAL_NOT;
  pointer_t<uint8_t> output_ptr(input.size());

  unary_transform(input_ptr, make_boolean_iterator(output_ptr), input.size(), op, no_cache, no_key);

  REQUIRE(std::vector<uint8_t>(output_ptr) == std::vector<uint8_t>{1, 0, 0, 0});
}

struct Transform_DifferentOutputTypes_Fixture_Tag;
C2H_TEST("Transform works with output of different type", "[transform]")
{
  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 24)));

  operation_t op               = make_operation("op",
                                  R"(struct pair { short a; size_t b; };
extern "C" __device__ void op(void* x_ptr, void* out_ptr) {
  int* x = static_cast<int*>(x_ptr);
  pair* out = static_cast<pair*>(out_ptr);
  *out = pair{ short(*x), size_t(*x) };
})");
  const std::vector<int> input = generate<int>(num_items);
  std::vector<pair> expected(num_items);
  std::vector<pair> output(num_items);
  for (std::size_t i = 0; i < num_items; ++i)
  {
    expected[i] = {short(input[i]), size_t(input[i])};
  }
  pointer_t<int> input_ptr(input);
  pointer_t<pair> output_ptr(output);

  auto& build_cache    = get_cache<Transform_DifferentOutputTypes_Fixture_Tag>();
  const auto& test_key = make_key<int, pair>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<pair>(output_ptr));
  }
}

struct alignas(8) unary_storage_in
{
  int x;
  short y;
};

struct alignas(16) unary_storage_out
{
  long long sum;
  int diff;

  bool operator==(const unary_storage_out& other) const
  {
    return sum == other.sum && diff == other.diff;
  }
};

struct Transform_UnaryStorageTypes_Fixture_Tag;
C2H_TEST("Transform works with unary storage types of different size/alignment", "[transform]")
{
  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));

  operation_t op = make_operation("op",
                                  R"(struct alignas(8) unary_storage_in { int x; short y; };
struct alignas(16) unary_storage_out { long long sum; int diff; };
extern "C" __device__ void op(void* x_ptr, void* out_ptr) {
  auto* x = static_cast<unary_storage_in*>(x_ptr);
  auto* out = static_cast<unary_storage_out*>(out_ptr);
  out->sum = static_cast<long long>(x->x) + x->y;
  out->diff = x->x - x->y;
})");

  std::vector<unary_storage_in> input(num_items);
  std::vector<unary_storage_out> output(num_items);
  std::vector<unary_storage_out> expected(num_items);
  for (std::size_t i = 0; i < num_items; ++i)
  {
    input[i]    = {static_cast<int>(i + 3), static_cast<short>(i % 7)};
    expected[i] = {static_cast<long long>(input[i].x) + input[i].y, input[i].x - input[i].y};
  }

  pointer_t<unary_storage_in> input_ptr(input);
  pointer_t<unary_storage_out> output_ptr(output);

  auto& build_cache    = get_cache<Transform_UnaryStorageTypes_Fixture_Tag>();
  const auto& test_key = make_key<unary_storage_in, unary_storage_out>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<unary_storage_out>(output_ptr));
  }
}

struct Transform_CustomTypes_Fixture_Tag;
C2H_TEST("Transform works with custom types", "[transform]")
{
  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 24)));

  operation_t op              = make_operation("op",
                                  R"(struct pair { short a; size_t b; };
extern "C" __device__ void op(void* x_ptr, void* out_ptr) {
  pair* x = static_cast<pair*>(x_ptr);
  pair* out = static_cast<pair*>(out_ptr);
  *out = pair{ x->a * 2, x->b * 2  };
})");
  const std::vector<short> a  = generate<short>(num_items);
  const std::vector<size_t> b = generate<size_t>(num_items);
  std::vector<pair> input(num_items);
  std::vector<pair> output(num_items);
  for (std::size_t i = 0; i < num_items; ++i)
  {
    input[i] = pair{a[i], b[i]};
  }
  pointer_t<pair> input_ptr(input);
  pointer_t<pair> output_ptr(output);

  auto& build_cache    = get_cache<Transform_CustomTypes_Fixture_Tag>();
  const auto& test_key = make_key<pair, pair>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<pair> expected(num_items, {0, 0});
  std::transform(input.begin(), input.end(), expected.begin(), [](const pair& x) {
    return pair{short(x.a * 2), x.b * 2};
  });
  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<pair>(output_ptr));
  }
}

struct Transform_CustomTypes_WellKnown_Fixture_Tag;
C2H_TEST("Transform works with custom types with well-known operators", "[transform][well_known]")
{
  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 24)));

  operation_t op_state = make_operation("op",
                                        R"(struct pair { short a; size_t b; };
extern "C" __device__ void op(void* x_ptr, void* out_ptr) {
  pair* x = static_cast<pair*>(x_ptr);
  pair* out = static_cast<pair*>(out_ptr);
  *out = pair{ x->a * 2, x->b * 2  };
})");
  cccl_op_t op         = op_state;
  // HACK: this doesn't actually match the operation above, but that's fine, as we are supposed to not take the
  // well-known path anyway
  op.type                     = cccl_op_kind_t::CCCL_NEGATE;
  const std::vector<short> a  = generate<short>(num_items);
  const std::vector<size_t> b = generate<size_t>(num_items);
  std::vector<pair> input(num_items);
  std::vector<pair> output(num_items);
  for (std::size_t i = 0; i < num_items; ++i)
  {
    input[i] = pair{a[i], b[i]};
  }
  pointer_t<pair> input_ptr(input);
  pointer_t<pair> output_ptr(output);

  auto& build_cache    = get_cache<Transform_CustomTypes_WellKnown_Fixture_Tag>();
  const auto& test_key = make_key<pair, pair>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<pair> expected(num_items, {0, 0});
  std::transform(input.begin(), input.end(), expected.begin(), [](const pair& x) {
    return pair{short(x.a * 2), x.b * 2};
  });
  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<pair>(output_ptr));
  }
}

struct Transform_InputIterators_Fixture_Tag;
C2H_TEST("Transform works with input iterators", "[transform]")
{
  const std::size_t num_items = GENERATE(1, 42, take(1, random(1 << 12, 1 << 16)));
  operation_t op              = make_operation("op", get_unary_op(get_type_info<int>().type));
  iterator_t<int, counting_iterator_state_t<int>> input_it = make_counting_iterator<int>("int");
  input_it.state.value                                     = 0;
  pointer_t<int> output_it(num_items);

  auto& build_cache    = get_cache<Transform_InputIterators_Fixture_Tag>();
  const auto& test_key = make_key<int>();

  unary_transform(input_it, output_it, num_items, op, build_cache, test_key);

  // vector storing a sequence of values 0, 1, 2, ..., num_items - 1
  std::vector<int> input(num_items);
  std::iota(input.begin(), input.end(), 0);

  std::vector<int> expected(num_items);
  std::transform(input.begin(), input.end(), expected.begin(), [](const int& x) {
    return x * 2;
  });
  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<int>(output_it));
  }
}

struct Transform_OutputIterators_Fixture_Tag;
C2H_TEST("Transform works with output iterators", "[transform]")
{
  const int num_items = GENERATE(1, 42, take(1, random(1 << 12, 1 << 16)));
  operation_t op      = make_operation("op", get_unary_op(get_type_info<int>().type));
  iterator_t<int, random_access_iterator_state_t<int>> output_it =
    make_random_access_iterator<int>(iterator_kind::OUTPUT, "int", "out", " * 2");
  const std::vector<int> input = generate<int>(num_items);
  pointer_t<int> input_it(input);
  pointer_t<int> inner_output_it(num_items);
  output_it.state.data = inner_output_it.ptr;

  auto& build_cache    = get_cache<Transform_OutputIterators_Fixture_Tag>();
  const auto& test_key = make_key<int>();

  unary_transform(input_it, output_it, num_items, op, build_cache, test_key);

  std::vector<int> expected(num_items);
  std::transform(input.begin(), input.end(), expected.begin(), [](int x) {
    return x * 4;
  });
  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<int>(inner_output_it));
  }
}

struct Transform_BinaryOp_Fixture_Tag;
C2H_TEST("Transform with binary operator", "[transform]")
{
  const std::size_t num_items   = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  const std::vector<int> input1 = generate<int>(num_items);
  const std::vector<int> input2 = generate<int>(num_items);
  const std::vector<int> output(num_items, 0);
  pointer_t<int> input1_ptr(input1);
  pointer_t<int> input2_ptr(input2);
  pointer_t<int> output_ptr(output);

  operation_t op = make_operation("op",
                                  R"(extern "C" __device__ void op(void* x_ptr, void* y_ptr, void* out_ptr  ) {
  int* x = static_cast<int*>(x_ptr);
  int* y = static_cast<int*>(y_ptr);
  int* out = static_cast<int*>(out_ptr);
  *out = (*x > *y) ? *x : *y;
})");

  auto& build_cache    = get_cache<Transform_BinaryOp_Fixture_Tag>();
  const auto& test_key = make_key<int>();

  binary_transform(input1_ptr, input2_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<int> expected(num_items, 0);
  std::transform(input1.begin(), input1.end(), input2.begin(), expected.begin(), [](const int& x, const int& y) {
    return (x > y) ? x : y;
  });

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<int>(output_ptr));
  }
}

struct alignas(16) binary_storage_in1
{
  long long a;
  int b;
};

struct alignas(8) binary_storage_in2
{
  int c;
  int d;
};

struct alignas(16) binary_storage_out
{
  long long sum;
  int diff;

  bool operator==(const binary_storage_out& other) const
  {
    return sum == other.sum && diff == other.diff;
  }
};

struct Transform_BinaryStorageTypes_Fixture_Tag;
C2H_TEST("Transform works with binary storage types of different size/alignment", "[transform]")
{
  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));

  operation_t op = make_operation("op",
                                  R"(struct alignas(16) binary_storage_in1 { long long a; int b; };
struct alignas(8) binary_storage_in2 { int c; int d; };
struct alignas(16) binary_storage_out { long long sum; int diff; };
extern "C" __device__ void op(void* x_ptr, void* y_ptr, void* out_ptr) {
  auto* x = static_cast<binary_storage_in1*>(x_ptr);
  auto* y = static_cast<binary_storage_in2*>(y_ptr);
  auto* out = static_cast<binary_storage_out*>(out_ptr);
  out->sum = x->a + static_cast<long long>(y->c);
  out->diff = x->b - y->d;
})");

  std::vector<binary_storage_in1> input1(num_items);
  std::vector<binary_storage_in2> input2(num_items);
  std::vector<binary_storage_out> output(num_items);
  std::vector<binary_storage_out> expected(num_items);
  for (std::size_t i = 0; i < num_items; ++i)
  {
    input1[i]   = {static_cast<long long>(i + 5), static_cast<int>(i + 2)};
    input2[i]   = {static_cast<int>(i + 7), static_cast<int>(i + 1)};
    expected[i] = {input1[i].a + static_cast<long long>(input2[i].c), input1[i].b - input2[i].d};
  }

  pointer_t<binary_storage_in1> input1_ptr(input1);
  pointer_t<binary_storage_in2> input2_ptr(input2);
  pointer_t<binary_storage_out> output_ptr(output);

  auto& build_cache    = get_cache<Transform_BinaryStorageTypes_Fixture_Tag>();
  const auto& test_key = make_key<binary_storage_in1, binary_storage_in2, binary_storage_out>();

  binary_transform(input1_ptr, input2_ptr, output_ptr, num_items, op, build_cache, test_key);

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<binary_storage_out>(output_ptr));
  }
}

struct Transform_BinaryOp_Iterator_Fixture_Tag;
C2H_TEST("Binary transform with one iterator", "[transform]")
{
  const std::size_t num_items   = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  const std::vector<int> input1 = generate<int>(num_items);

  iterator_t<int, counting_iterator_state_t<int>> input2_it = make_counting_iterator<int>("int");
  input2_it.state.value                                     = 0;

  const std::vector<int> output(num_items, 0);
  pointer_t<int> input1_ptr(input1);
  pointer_t<int> output_ptr(output);

  operation_t op = make_operation("op",
                                  R"(extern "C" __device__ void op(void* x_ptr, void* y_ptr, void* out_ptr) {
  int* x = static_cast<int*>(x_ptr);
  int* y = static_cast<int*>(y_ptr);
  int* out = static_cast<int*>(out_ptr);
  *out = (*x > *y) ? *x : *y;
})");

  auto& build_cache    = get_cache<Transform_BinaryOp_Iterator_Fixture_Tag>();
  const auto& test_key = make_key<int>();

  binary_transform(input1_ptr, input2_it, output_ptr, num_items, op, build_cache, test_key);

  std::vector<int> input2(num_items);
  std::iota(input2.begin(), input2.end(), 0);
  std::vector<int> expected(num_items, 0);
  std::transform(input1.begin(), input1.end(), input2.begin(), expected.begin(), [](const int& x, const int& y) {
    return (x > y) ? x : y;
  });

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<int>(output_ptr));
  }
}

using floating_point_types = c2h::type_list<
#if _CCCL_HAS_NVFP16()
  __half,
#endif
  float,
  double>;
struct Transform_FloatingPointTypes_Fixture_Tag;
C2H_TEST("Transform works with floating point types", "[transform]", floating_point_types)
{
  using T = c2h::get<0, TestType>;

  const std::size_t num_items      = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  operation_t op                   = make_operation("op", get_unary_op(get_type_info<T>().type));
  const std::vector<int> int_input = generate<int>(num_items);
  // Suppress harmless conversion warnings on MSVC
  _CCCL_DIAG_PUSH
  _CCCL_DIAG_SUPPRESS_MSVC(4244)
  const std::vector<T> input(int_input.begin(), int_input.end());
  _CCCL_DIAG_POP
  const std::vector<T> output(num_items, 0);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr(output);

  auto& build_cache    = get_cache<Transform_FloatingPointTypes_Fixture_Tag>();
  const auto& test_key = make_key<T>();

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<T> expected(num_items, 0);
  std::transform(input.begin(), input.end(), expected.begin(), [](const T& x) {
    return T{2} * x;
  });

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<T>(output_ptr));
  }
}

C2H_TEST("Transform works with C++ source operations", "[transform]")
{
  using T = int32_t;

  const std::size_t num_items = GENERATE(42, 1337, 42000);

  // Create operation from C++ source instead of LTO-IR
  std::string cpp_source = R"(
    extern "C" __device__ void op(void* input, void* output) {
      int* in = (int*)input;
      int* out = (int*)output;
      *out = *in * 2;
    }
  )";

  operation_t op = make_cpp_operation("op", cpp_source);

  const std::vector<T> input = generate<T>(num_items);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr(num_items);

  // Test key including flag that this uses C++ source
  std::optional<std::string> test_key = std::format("cpp_source_test_{}_{}", num_items, typeid(T).name());

  auto& cache = fixture<transform_build_cache_t, Transform_IntegralTypes_Fixture_Tag>::get_or_create().get_value();
  std::optional<transform_build_cache_t> cache_opt = cache;

  unary_transform(input_ptr, output_ptr, num_items, op, cache_opt, test_key);

  const std::vector<T> output = output_ptr;
  std::vector<T> expected     = input;
  std::transform(expected.begin(), expected.end(), expected.begin(), [](T x) {
    return x * 2;
  });
  REQUIRE(output == expected);
}

C2H_TEST("Transform works with C++ source operations using custom headers", "[transform]")
{
  using T = int32_t;

  const std::size_t num_items = GENERATE(42, 1337, 42000);

  // Create operation from C++ source that uses the identity function from header
  std::string cpp_source = R"(
    #include "test_identity.h"
    extern "C" __device__ void op(void* input, void* output) {
      int* in = (int*)input;
      int* out = (int*)output;
      int val = test_identity(*in);
      *out = val * 2;
    }
  )";

  operation_t op = make_cpp_operation("op", cpp_source);

  const std::vector<T> input = generate<T>(num_items);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr(num_items);

  // Test _ex version with custom build configuration
  const char* extra_flags[] = {"-DTEST_IDENTITY_ENABLED"};
  const char* extra_dirs[]  = {TEST_INCLUDE_PATH};
  cccl_build_config config  = make_build_config(extra_flags, 1, extra_dirs, 1);

  // Build with _ex version
  cccl_device_transform_build_result_t build{};
  const auto& build_info = BuildInformation<>::init();
  REQUIRE(
    CUDA_SUCCESS
    == cccl_device_unary_transform_build_ex(
      &build,
      input_ptr,
      output_ptr,
      op,
      build_info.get_cc_major(),
      build_info.get_cc_minor(),
      build_info.get_cub_path(),
      build_info.get_thrust_path(),
      build_info.get_libcudacxx_path(),
      build_info.get_ctk_path(),
      &config));

  // Execute the transform
  REQUIRE(CUDA_SUCCESS == cccl_device_unary_transform(build, input_ptr, output_ptr, num_items, op, CU_STREAM_LEGACY));

  // Verify results
  std::vector<T> output(num_items);
  cudaMemcpy(output.data(), static_cast<void*>(output_ptr.ptr), sizeof(T) * num_items, cudaMemcpyDeviceToHost);
  std::vector<T> expected = input;
  std::transform(expected.begin(), expected.end(), expected.begin(), [](T x) {
    return x * 2;
  });
  REQUIRE(output == expected);

  // Cleanup
  REQUIRE(CUDA_SUCCESS == cccl_device_transform_cleanup(&build));
}

struct transform_stateful_counter_state_t
{
  int* d_counter;
};

C2H_TEST("Transform works with stateful unary operators", "[transform]")
{
  const std::size_t num_items = GENERATE(0, 42, take(4, random(1 << 12, 1 << 16)));
  const std::vector<int> host_counter{0};
  pointer_t<int> counter(host_counter);
  stateful_operation_t<transform_stateful_counter_state_t> op = make_operation(
    "op",
    R"(struct transform_stateful_counter_state_t { int* d_counter; };
extern "C" __device__ void op(void* state_ptr, void* x_ptr, void* out_ptr) {
  auto* state = static_cast<transform_stateful_counter_state_t*>(state_ptr);
  atomicAdd(state->d_counter, 1);
  int x = *static_cast<int*>(x_ptr);
  *static_cast<int*>(out_ptr) = x * 2;
})",
    transform_stateful_counter_state_t{counter.ptr});

  const std::vector<int> input = generate<int>(num_items);
  const std::vector<int> output(num_items, 0);
  pointer_t<int> input_ptr(input);
  pointer_t<int> output_ptr(output);

  std::optional<transform_build_cache_t> build_cache = std::nullopt;
  std::optional<std::string> test_key                = std::nullopt;

  unary_transform(input_ptr, output_ptr, num_items, op, build_cache, test_key);

  std::vector<int> expected(num_items, 0);
  std::transform(input.begin(), input.end(), expected.begin(), [](int x) {
    return x * 2;
  });

  if (num_items > 0)
  {
    REQUIRE(expected == std::vector<int>(output_ptr));
    REQUIRE(counter[0] == static_cast<int>(num_items));
  }
}

#ifndef CCCL_C_PARALLEL_V2
C2H_TEST("Transform build result has serialization metadata populated", "[transform][serialization]")
{
  using T = int32_t;

  constexpr int device_id = 0;
  const auto& build_info  = BuildInformation<device_id>::init();

  cccl_op_t op = make_well_known_unary_operation();
  pointer_t<T> in(1);
  pointer_t<T> out(1);

  BuildResultT build{};
  REQUIRE(
    CUDA_SUCCESS
    == cccl_device_unary_transform_build(
      &build,
      in,
      out,
      op,
      build_info.get_cc_major(),
      build_info.get_cc_minor(),
      build_info.get_cub_path(),
      build_info.get_thrust_path(),
      build_info.get_libcudacxx_path(),
      build_info.get_ctk_path()));

  CHECK(build.cc == build_info.get_cc_major() * 10 + build_info.get_cc_minor());
  CHECK((build.payload != nullptr && build.payload_kind == CCCL_PAYLOAD_CUBIN));
  CHECK(build.payload_size > 0);
  CHECK(build.runtime_policy != nullptr);
  CHECK(build.runtime_policy_size > 0);
  REQUIRE(build.transform_kernel_lowered_name != nullptr);
  CHECK(build.transform_kernel_lowered_name[0] != '\0');

  REQUIRE(CUDA_SUCCESS == cccl_device_transform_cleanup(&build));
}

C2H_TEST("Transform compile/load round-trip", "[transform][serialization]")
{
  using T = int32_t;

  constexpr int device_id = 0;
  const auto& build_info  = BuildInformation<device_id>::init();

  cccl_op_t op = make_well_known_unary_operation();
  pointer_t<T> dummy_in(1);
  pointer_t<T> dummy_out(1);

  BuildResultT build{};
  REQUIRE(
    CUDA_SUCCESS
    == cccl_device_unary_transform_compile(
      &build,
      dummy_in,
      dummy_out,
      op,
      build_info.get_cc_major(),
      build_info.get_cc_minor(),
      build_info.get_cub_path(),
      build_info.get_thrust_path(),
      build_info.get_libcudacxx_path(),
      build_info.get_ctk_path(),
      nullptr));

  REQUIRE((build.payload != nullptr && build.payload_kind == CCCL_PAYLOAD_CUBIN));
  REQUIRE(build.payload_size > 0);
  REQUIRE(build.transform_kernel_lowered_name != nullptr);
  CHECK(build.library == nullptr);
  CHECK(build.transform_kernel == nullptr);

  REQUIRE(CUDA_SUCCESS == cccl_device_transform_load(&build));
  REQUIRE(build.library != nullptr);
  CHECK(build.transform_kernel != nullptr);

  constexpr std::size_t n    = 16;
  const std::vector<T> input = generate<T>(n);
  pointer_t<T> input_ptr(input);
  pointer_t<T> output_ptr(n);
  CUstream null_stream = nullptr;

  REQUIRE(CUDA_SUCCESS == cccl_device_unary_transform(build, input_ptr, output_ptr, n, op, null_stream));

  std::vector<T> expected(input);
  std::transform(expected.begin(), expected.end(), expected.begin(), [](T x) {
    return -x;
  });
  REQUIRE(expected == std::vector<T>(output_ptr));

  REQUIRE(CUDA_SUCCESS == cccl_device_transform_cleanup(&build));
}
#endif // CCCL_C_PARALLEL_V2