#include <pls/pls.h>
#include <pls/internal/helpers/profiler.h>
#include <pls/internal/helpers/mini_benchmark.h>

#include "node.h"

const int SEED = 42;
const int ROOT_CHILDREN = 140;
const double Q = 0.124875;
const int NORMAL_CHILDREN = 8;

const int NUM_NODES = 71069;

int count_child_nodes(uts::node &node) {
  int child_count = 1;
  std::vector<uts::node> children = node.spawn_child_nodes();

  if (children.empty()) {
    return child_count;
  }

  std::vector<int> results(children.size());
  for (size_t i = 0; i < children.size(); i++) {
    size_t index = i;
    auto lambda = [&, index] { results[index] = count_child_nodes(children[index]); };
    pls::lambda_task_by_value<typeof(lambda)> sub_task(lambda);
    pls::scheduler::spawn_child(sub_task);
  }
  pls::scheduler::wait_for_all();
  for (auto result : results) {
    child_count += result;
  }

  return child_count;
}

int unbalanced_tree_search(int seed, int root_children, double q, int normal_children) {
  int result;

  auto lambda = [&] {
    uts::node root(seed, root_children, q, normal_children);
    result = count_child_nodes(root);
  };
  pls::lambda_task_by_reference<typeof(lambda)> sub_task(lambda);
  pls::scheduler::spawn_child(sub_task);
  pls::scheduler::wait_for_all();

  return result;
}

int main() {
  PROFILE_ENABLE
  pls::internal::helpers::run_mini_benchmark([&] {
    unbalanced_tree_search(SEED, ROOT_CHILDREN, Q, NORMAL_CHILDREN);
  }, 8, 2000);

  PROFILE_SAVE("test_profile.prof")
}

//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")
//}