#ifndef SCHEDULERS_HPP
#define SCHEDULERS_HPP

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

using namespace std;

class ProlificScheduler {
public:
	static void execute(int* shared_tensor, int n, int k) {
		auto t_start = chrono::high_resolution_clock::now();
		pid_t root = getpid();

		for (int i = 0; i<k;i++) {
			if (getpid() != root) {
				break;
			}

			pid_t pid = fork();

			if (pid == 0) {
				for (int row = 0; row < n;row++) {
					int* start_of_row = shared_tensor + (i*n*n) + (row*n);
					int* end_of_row = start_of_row + n;
					std::sort(start_of_row, end_of_row);
				}
				_exit(0);
				
			}
		}

		
		auto t_setup = chrono::high_resolution_clock::now();

		if (getpid() == root) {
			int status;
			while (waitpid(-1, &status, 0) > 0);

			auto t_end = chrono::high_resolution_clock::now();

			chrono::duration<double, std::milli> setup_ms = t_setup - t_start;
			chrono::duration<double, std::milli> exec_ms  = t_end - t_setup;

			cout << " -> Prolific Setup Time:  " << setup_ms.count() << "ms\n";
			cout << " -> Prolific Execution Time:  " << exec_ms.count() << "ms\n";
		}
	}
};



class CollectiveScheduler {
public:
	static void execute(int* shared_tensor, int n, int k) {
		auto t_start = chrono::high_resolution_clock::now();
		pid_t root = getpid();

		int levels = std::ceil(std::log2(k));
		if (levels == 0) levels = 1;

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

		auto t_setup = chrono::high_resolution_clock::now();

		if (getpid() != root) {
			int total_workers = (1 << levels);
			int worker_id = getpid() % total_workers;

			for (int i = 0;i < k;i++) {
				if (i % total_workers == worker_id % total_workers) {
					for (int row = 0; row < n; row++) {
						int* start_of_row = shared_tensor + (i*n*n) + (row*n);
						int* end_of_row = start_of_row + n;
						std::sort(start_of_row, end_of_row);
					}
				}
			}

			_exit(0);

		}

		if (getpid() == root) {
			int status;
			while (waitpid(-1, &status, 0) > 0);

			auto t_end = chrono::high_resolution_clock::now();

			chrono::duration<double, std::milli> setup_ms = t_setup - t_start;
			chrono::duration<double, std::milli> exec_ms = t_end - t_setup;

			cout  << " -> Collective Setup Time:  " << setup_ms.count() << "ms\n";
			cout << " -> Collective Execution Time: "<< exec_ms.count() << "ms\n";
		}
	}
};

#endif