#include <iostream>
#include <iomanip>
#include <cstdlib>
#include <ctime>
#include <sys/mman.h>
#include "schedulers.hpp"

using namespace std;

void print_tensor(int* tensor, int n,int k, const string& title) {
	cout << "\n=== " << title << " ===\n";
	for(int page = 0; page < k;page++) {
		cout << "--- Page (k=" << page + 1 << ") ---\n";

		for(int i = 0;i<n;i++) {
			for(int j = 0;j<n;j++) {
				cout << setw(6) << tensor[(page*n*n) + (i*n) + j] << " ";
			}
			cout <<"\n";
		}
	}
	cout << "======================\n";
}

void fill_random(int* tensor, int size) {
	for (int i = 0;i < size; i++) {
		tensor[i] =  (rand() % 2001) - 1000;
	}
}


int main() {
	srand(time(NULL));

	int n = 4;

	int test_k[] = {10, 100};

	for (int k : test_k) {
		cout << "\n*************************************************\n";
		cout << "TESTING FOR n = " << n << " AND k =  " << k << " WORKERS";
		cout << "\n*************************************************\n";

		int total_elements = n * n *k;
	


		int* B = (int*)mmap(NULL, total_elements * sizeof(int), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);

		if(B == MAP_FAILED) {
			cerr << "Memory allocation failed!\n";
			return 1;
		}

		//Printing Tensors with Prolific Scheduler
		fill_random(B, total_elements);
		if (k == 10) print_tensor(B, n, k, "Original Tensor B (Prolific Test)");

		cout << "\n[Running Prolififc Scheduler...]\n";
		ProlificScheduler::execute(B, n, k);

		if (k == 10) print_tensor(B, n, k, "Sorted Tensor B (Prolific)");

		//Printing Tensors with Collective Scheduler
		fill_random(B, total_elements);

		cout << "\n[Running Collective Scheduler...]\n";
		CollectiveScheduler::execute(B, n, k);

		if (k == 10) print_tensor(B, n, k, "Sorted Tensor B (Collective)");

		munmap(B, total_elements * sizeof(int));
	}

	return 0;
}