#ifndef SCHEDULERS_H
#define SCHEDULERS_H

#include <iostream>
#include <unistd.h>
#include <sys/wait.h>
#include <vector>
#include <algorithm>
#include <cmath>
#include <functional>

//Class 1: PROLIFIC SCHEDULER
class ProlificScheduler {
public:
    static void execute(int levels, int total_tasks, std::function<void(int)> task_func) {
	pid_t root = getpid();
	std::vector<pid_t> pids;
	int worker_id = -1;

	for (int i =0; i < levels; ++i) {
	    if (getpid() != root) { break;}
	    pid_t pid = fork () ;
	    if (pid == 0) {
	       worker_id = i;
	       break;
	    } else {
	       pids.push_back(pid);
	    }
	}

	if (getpid() == root) {
	    int status;
	    while (true) {
	        pid_t finished = waitpid(-1, &status, 0);
	        if (finished > 0) {
                    pids.erase(std::remove(pids.begin(), pids.end(), finished), pids.end());
                    continue;
                }
                if (finished == -1 && errno == ECHILD) break;
                break;
           }
        } else {
            for (int i =0; i < total_tasks; ++i) {
                if (i % levels == worker_id) {
                    task_func(i); //Execute the sorting task
                }
            }
            _exit(0);
        }
     }
};

//Class 2: COLLECTIVE SCHEDULER
class CollectiveScheduler {
public:
    static void execute(int levels, int total_tasks, std::function<void(int)> task_func) {
        pid_t root = getpid();

        for (int i =0; i <levels; ++i) {
            fork();
        }

        if (getpid() != root) {
            int worker_id = getpid() % root;
            int worker_count = std::pow(2, levels) -1;

            for (int i =0; i < total_tasks; ++i) {
                if (i %  worker_count == worker_id % worker_count) {
                    task_func(i); //Execute the sorting task
                }
            }
            _exit(0);
        }

        int status;
        while (true) {
            pid_t finished = waitpid(-1, &status, 0);
            if (finished > 0) continue;
            if (finished == -1 && errno == ECHILD) break;
       }
    }
};
#endif