#include "pls/internal/scheduling/scheduler.h"
#include "pls/internal/scheduling/static_scheduler_memory.h"

using namespace pls::internal::scheduling;

#include "benchmark_runner.h"
#include "benchmark_base/fft.h"

using namespace comparison_benchmarks::base;

void pls_conquer(fft::complex_vector::iterator data, fft::complex_vector::iterator swap_array, int n) {
  if (n < 2) {
    return;
  }

  fft::divide(data, swap_array, n);
  if (n <= fft::RECURSIVE_CUTOFF) {
    fft::conquer(data, swap_array, n / 2);
    fft::conquer(data + n / 2, swap_array + n / 2, n / 2);
  } else {
    scheduler::spawn([data, n, swap_array]() {
      pls_conquer(data, swap_array, n / 2);
    });
    scheduler::spawn([data, n, swap_array]() {
      pls_conquer(data + n / 2, swap_array + n / 2, n / 2);
    });
    scheduler::sync();
  }

  fft::combine(data, n);
}

constexpr int MAX_NUM_THREADS = 8;
constexpr int MAX_NUM_TASKS = 32;
constexpr int MAX_STACK_SIZE = 1024 * 8;

int main(int argc, char **argv) {
  int num_threads;
  string directory;
  benchmark_runner::read_args(argc, argv, num_threads, directory);

  string test_name = to_string(num_threads) + ".csv";
  string full_directory = directory + "/PLS_v3/";
  benchmark_runner runner{full_directory, test_name};

  fft::complex_vector data(fft::SIZE);
  fft::complex_vector swap_array(fft::SIZE);
  fft::fill_input(data);

  static_scheduler_memory<MAX_NUM_THREADS,
                          MAX_NUM_TASKS,
                          MAX_STACK_SIZE> global_scheduler_memory;
  scheduler scheduler{global_scheduler_memory, (unsigned) num_threads};

  runner.run_iterations(fft::NUM_ITERATIONS, [&]() {
    scheduler.perform_work([&]() {
      pls_conquer(data.begin(), swap_array.begin(), fft::SIZE);;
    });
  }, fft::NUM_WARMUP_ITERATIONS, [&]() {
    fft::fill_input(data); // Reset data before each run
  });
  runner.commit_results(true);

  return 0;
}
