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

using namespace pls::internal::scheduling;

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

using namespace comparison_benchmarks::base;

#include <atomic>

int count_child_nodes(unbalanced::node &node) {
  if (node.get_num_children() < 1) {
    return 1;
  }

  std::atomic<int> count{1};
  for (int i = 0; i < node.get_num_children(); i++) {
    scheduler::spawn([i, &count, &node] {
      unbalanced::node child_node = node.spawn_child_node(i);
      count.fetch_add(count_child_nodes(child_node));
    });
  }
  scheduler::sync();

  return count;
}

int unbalanced_tree_search(int seed, int root_children, double q, int normal_children) {
  unbalanced::node root(seed, root_children, q, normal_children);
  return count_child_nodes(root);
}

constexpr int MAX_NUM_TASKS = 256;
constexpr int MAX_STACK_SIZE = 1024 * 2;

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};
  runner.enable_memory_stats();

  scheduler scheduler{(unsigned) num_threads, MAX_NUM_TASKS, MAX_STACK_SIZE};

//  scheduler.get_profiler().disable_memory_measure();
  runner.run_iterations(1, [&]() {
    scheduler.perform_work([&]() {
      unbalanced_tree_search(unbalanced::SEED,
                             unbalanced::ROOT_CHILDREN,
                             unbalanced::Q,
                             unbalanced::NORMAL_CHILDREN);
    });
  }, 0);
//  scheduler.get_profiler().current_run().print_dag(std::cout);
//  scheduler.get_profiler().current_run().print_stats();
  runner.commit_results(true);

}

//int main() {
//  PROFILE_ENABLE
//  pls::malloc_scheduler_memory my_scheduler_memory{8, 2u << 18};
//  pls::scheduler scheduler{&my_scheduler_memory, 8};
//
//  scheduler.perform_work([&] {
//    PROFILE_MAIN_THREAD
//    for (int i = 0; i < 50; i++) {
//      PROFILE_WORK_BLOCK("Top Level")
//      int result = unbalanced_tree_search(SEED, ROOT_CHILDREN, Q, NORMAL_CHILDREN);
//      std::cout << result << std::endl;
//    }
//  });
//
//  PROFILE_SAVE("test_profile.prof")
//}