#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;
  }

  auto current_task = pls::task::current();
  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);
    current_task->spawn_child(sub_task);
  }
  current_task->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) {
  static auto id = pls::unique_id::create(42);
  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)> task(lambda);
  pls::lambda_task_by_reference<typeof(lambda)> sub_task(lambda);
  pls::task root_task{&sub_task, id};
  pls::scheduler::execute_task(root_task);

  return result;
}
//
//int main() {
//  PROFILE_ENABLE
//  pls::internal::helpers::run_mini_benchmark([&] {
//    unbalanced_tree_search(SEED, ROOT_CHILDREN, Q, NORMAL_CHILDREN);
//  }, 8, 4000);
//
//  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")
}