// This program demonstrates how to change the worker behavior
// upon the creation of an executor.

#include <taskflow/taskflow.hpp>


// ----------------------------------------------------------------------------
// Affinity
// ----------------------------------------------------------------------------
#if defined(__linux__)
 #include <sched.h>
 #include <pthread.h>
#elif defined(_WIN32)
 #include <windows.h>
#elif defined(__APPLE__)
 #include <mach/mach.h>
 #include <mach/thread_policy.h>
#endif

// affine the given thread to a specific core
bool affine(std::thread& thread, size_t core_id) {
#if defined(__linux__)
 cpu_set_t cpuset;
 CPU_ZERO(&cpuset);
 CPU_SET(core_id, &cpuset);
 pthread_t native_handle = thread.native_handle();
 return pthread_setaffinity_np(native_handle, sizeof(cpu_set_t), &cpuset) == 0;
#elif defined(_WIN32)
 return SetThreadAffinityMask(thread.native_handle(), 1ULL << core_id) != 0;
#elif defined(__APPLE__)
 thread_port_t native_handle = pthread_mach_thread_np(thread.native_handle());
 thread_affinity_policy_data_t policy = {static_cast<integer_t>(core_id)};
 return thread_policy_set(
 native_handle, THREAD_AFFINITY_POLICY, (thread_policy_t)&policy, 1
 ) == KERN_SUCCESS;
#else
 // Unsupported platform
 return false;
#endif
}

// ----------------------------------------------------------------------------

class CustomWorkerBehavior : public tf::WorkerInterface {

 public:
 
 // to call before the worker enters the scheduling loop
 void scheduler_prologue(tf::Worker& w) override {
 printf("worker %zu prepares to enter the work-stealing loop\n", w.id());
 
 // now affine the worker to a particular CPU core equal to its id
 if(affine(w.thread(), w.id())) {
 printf("successfully affines worker %zu to CPU core %zu\n", w.id(), w.id());
 }
 else {
 printf("failed to affine worker %zu to CPU core %zu\n", w.id(), w.id());
 }
 }

 // to call after the worker leaves the scheduling loop
 void scheduler_epilogue(tf::Worker& w, std::exception_ptr) override {
 printf("worker %zu left the work-stealing loop\n", w.id());
 }
};

int main() {
 tf::Executor executor(4, tf::make_worker_interface<CustomWorkerBehavior>());
 return 0;
}