
#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 parallel_scan(InIter input_start, const InIter input_end, OutIter output, BinaryOp op, Type neutral_element) {
  constexpr auto chunks_per_thread = 4;
  using namespace pls::internal::scheduling;

  auto size = std::distance(input_start, input_end);
  auto my_state = thread_state::get();
  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);

  // First Pass -> each chunk isolated
  // TODO: Put this onto the stack, not into heap allocated memory
  Type *tmp = new Type[chunks]; //reinterpret_cast<Type *>(my_state->task_stack_->push_bytes(sizeof(Type) * (chunks)));
  parallel_for(0, chunks,
               [tmp, input_start, input_end, output, items_per_chunk, &op, &neutral_element](int i) {
                 auto local_start = input_start + items_per_chunk * i;
                 auto local_end = local_start + items_per_chunk;
                 if (local_end > input_end) {
                   local_end = input_end;
                 }

                 // TODO: This MUST be dynamic to make sense, as it has a far bigger influence than any other cutoff
                 //
                 internal::serial_scan(local_start, local_end, output + items_per_chunk * i, op, neutral_element);
                 tmp[i] = *(output + std::distance(local_start, local_end) - 1);
               });
  internal::serial_scan(tmp, tmp + chunks, tmp, std::plus<int>(), 0);

  // Sum up pivot elements...
  auto output_start = output;
  auto output_end = output + size;
  parallel_for(1, chunks,
               [tmp, output_start, output_end, items_per_chunk, &op, &neutral_element](int i) {
                 auto local_start = output_start + items_per_chunk * i;
                 auto local_end = local_start + items_per_chunk;
                 if (local_end > output_end) {
                   local_end = output_end;
                 }

                 for (; local_start != local_end; local_start++) {
                   *local_start = op(*local_start, tmp[i - 1]);
                 }
               });

//  scheduler::wait_for_all();
  delete[] tmp;
}

}
}

#endif //PLS_PARALLEL_SCAN_IMPL_H_
