// filename: inclusive_scan.cu
#include <cuda/std/numeric>

#include <iostream>
#include <tuple>
#include <vector>

template <typename OffsetT>
std::tuple<std::vector<OffsetT>, std::vector<OffsetT>> make_in_out_offsets(const std::vector<OffsetT> sizes, OffsetT gap)
{
  std::vector<OffsetT> offsets;

  std::size_t segment_count = sizes.size();

  static constexpr OffsetT zero{0};

  offsets.resize(segment_count + 1);
  offsets[0] = zero;

  // cuda::std::plus<> plus_t{};

  cuda::std::inclusive_scan(sizes.begin(), sizes.end(), offsets.begin() + 1);

  std::vector<OffsetT> sizes_with_gaps;
  sizes_with_gaps.resize(segment_count);
  for (std::size_t i = 0; i < segment_count; ++i)
  {
    const auto s       = sizes[i];
    sizes_with_gaps[i] = (s == 0) ? gap : s;
  }

  std::vector<OffsetT> offsets_with_gaps;
  offsets_with_gaps.resize(segment_count + 1);
  offsets_with_gaps[0] = zero;
  cuda::std::inclusive_scan(sizes_with_gaps.begin(), sizes_with_gaps.end(), offsets_with_gaps.begin() + 1);

  return {offsets, offsets_with_gaps};
}

int main()
{
  std::vector<unsigned int> sizes = {2, 4, 0, 5, 1, 0, 0, 1};

  auto [a, b] = make_in_out_offsets(sizes, 4u);

  for (auto el : b)
  {
    std::cout << el << " ";
  }

  std::cout << "\n";

  return 0;
}