#include "pls/internal/scheduling/scheduler.h"
#include "pls/internal/scheduling/parallel_result.h"
#include "pls/internal/scheduling/scheduler_memory.h"
using namespace pls::internal::scheduling;

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

parallel_result<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]);
    };
    using child_type = pls::lambda_task_by_value<typeof(lambda)>;
    pls::scheduler::spawn_child<child_type>(lambda);
  }
  pls::scheduler::wait_for_all();
  for (auto result : results) {
    child_count += result;
  }

  return child_count;
}

parallel_result<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);
  };
  using child_type = pls::lambda_task_by_reference<typeof(lambda)>;
  pls::scheduler::spawn_child<child_type>(lambda);
  pls::scheduler::wait_for_all();

  return result;
}

constexpr size_t MAX_NUM_THREADS = 5;

constexpr size_t MAX_NUM_TASKS = 128;

constexpr size_t MAX_NUM_CONTS = 128;
constexpr size_t MAX_CONT_SIZE = 512;

volatile int result;
int main() {
  PROFILE_ENABLE
  static_scheduler_memory<MAX_NUM_THREADS,
                          MAX_NUM_TASKS,
                          MAX_NUM_CONTS,
                          MAX_CONT_SIZE> static_scheduler_memory;

  scheduler scheduler{static_scheduler_memory, MAX_NUM_THREADS};

  scheduler.perform_work([&]() {
    return scheduler::par([&]() {
      return unbalanced_tree_search(SEED, ROOT_CHILDREN, Q, NORMAL_CHILDREN);
    }, []() {
      return parallel_result<int>{0};
    }).then([](int a, int) {
      result = a;
      return parallel_result<int>{0};
    });
  });

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