
#ifndef PLS_PARALLEL_SCAN_IMPL_H_
#define PLS_PARALLEL_SCAN_IMPL_H_

#include <memory>
#include <functional>

#include "pls/pls.h"
#include "pls/internal/scheduling/thread_state.h"

namespace pls {
namespace algorithm {
namespace internal {
template<typename InIter, typename OutIter, typename BinaryOp, typename Type>
void serial_scan(InIter input_start, const InIter input_end, OutIter output, BinaryOp op, Type neutral_element) {
  auto current_input = input_start;
  auto current_output = output;
  auto last_value = neutral_element;
  while (current_input != input_end) {
    last_value = op(last_value, *current_input);
    *current_output = last_value;

    current_input++;
    current_output++;
  }
}

}

template<typename InIter, typename OutIter, typename BinaryOp, typename Type>
void scan(InIter in_start, const InIter in_end, OutIter out, BinaryOp op, Type neutral_elem) {
  constexpr auto chunks_per_thread = 4;
  using namespace pls::internal::scheduling;

  // TODO: This must be dynamic to make sense, as it has a far bigger influence than any other cutoff.
  //       The current strategy is static partitioning, and suboptimal in inballanced workloads.
  auto size = std::distance(in_start, in_end);
  auto num_threads = thread_state::get()->scheduler_->num_threads();
  auto chunks = num_threads * chunks_per_thread;
  auto items_per_chunk = std::max(1l, size / chunks);

  scheduler::allocate_on_stack(sizeof(Type) * (chunks), [&](void *memory) {
    Type *chunk_sums = reinterpret_cast<Type *>(memory);

    // First Pass = calculate each chunks individual prefix sum
    for_each_range(0, chunks, [&](int i) {
      auto chunk_start = in_start + items_per_chunk * i;
      auto chunk_end = std::min(in_end, chunk_start + items_per_chunk);
      auto chunk_output = out + items_per_chunk * i;

      internal::serial_scan(chunk_start, chunk_end, chunk_output, op, neutral_elem);
      chunk_sums[i] = *(out + std::distance(chunk_start, chunk_end) - 1);
    });

    // Calculate prefix sums of each chunks sum
    // (effectively the prefix sum at the end of each chunk, then used to correct the following chunk).
    internal::serial_scan(chunk_sums, chunk_sums + chunks, chunk_sums, std::plus<int>(), 0);

    // Second Pass = Use results from first pass to correct each chunks sum
    auto output_start = out;
    auto output_end = out + size;
    for_each_range(1, chunks, [&](int i) {
      auto chunk_start = output_start + items_per_chunk * i;
      auto chunk_end = std::min(output_end, chunk_start + items_per_chunk);

      for (; chunk_start != chunk_end; chunk_start++) {
        *chunk_start = op(*chunk_start, chunk_sums[i - 1]);
      }
    });
  });

  // End this work section by cleaning up stack and tasks
  scheduler::wait_for_all();
}

}
}

#endif //PLS_PARALLEL_SCAN_IMPL_H_
