#include "pls/pls.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++) {
    pls::spawn([i, &count, &node] {
      unbalanced::node child_node = node.spawn_child_node(i);
      count.fetch_add(count_child_nodes(child_node));
    });
  }
  pls::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 = 180;
constexpr int MAX_STACK_SIZE = 4096 * 1;

int main(int argc, char **argv) {
  auto settings = benchmark_runner::parse_parameters(argc, argv);

  string test_name = to_string(settings.num_threads_) + ".csv";
  string full_directory = settings.output_directory_ + "/PLS_v3/";
  benchmark_runner runner{full_directory, test_name};

  pls::scheduler scheduler{(unsigned) settings.num_threads_, MAX_NUM_TASKS, MAX_STACK_SIZE};

  if (settings.type_ == benchmark_runner::benchmark_settings::ISOLATED) {
#if PLS_PROFILING_ENABLED
    scheduler.get_profiler().disable_memory_measure();
    runner.add_custom_stats_field("T_1");
    runner.add_custom_stats_field("T_inf");
#endif
    printf("Running isolated measurement...\n");
    runner.enable_memory_stats();
    runner.pre_allocate_stats();

    runner.run_iterations(settings.iterations_, [&]() {
      scheduler.perform_work([&]() {
        unbalanced_tree_search(unbalanced::SEED,
                               unbalanced::ROOT_CHILDREN,
                               unbalanced::Q,
                               unbalanced::NORMAL_CHILDREN);
      });
    }, [&]() {}, [&]() {
#if PLS_PROFILING_ENABLED
      runner.store_custom_stat("T_1", scheduler.get_profiler().current_run().t_1_);
      runner.store_custom_stat("T_inf", scheduler.get_profiler().current_run().t_inf_);
#endif
    });
    runner.commit_results(true);
  } else {
    printf("Running periodic measurement...\n");
    runner.enable_wall_time_stats();
    runner.pre_allocate_stats();

    runner.run_periodic(settings.iterations_, settings.interval_period_, settings.interval_deadline_, [&]() {
      scheduler.perform_work([&]() {
        unbalanced_tree_search(unbalanced::SEED,
                               unbalanced::ROOT_CHILDREN,
                               unbalanced::Q,
                               unbalanced::NORMAL_CHILDREN);
      });
    });
    runner.commit_results(true);
  }

}