#include <iostream>
#include <vector>
#include <random>
#include <string>
#include <cstdlib>
#include "prolific_scheduler.hpp"
#include "collective_scheduler.hpp"

int main(int argc, char* argv[]) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " n {prolific|collective}\n";
        return 1;
    }

    int n = std::atoi(argv[1]);
    std::string mode = argv[2];

    //generates random matrix A(n,n) with values in [-1000, 1000]
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dist(-1000, 1000);

    std::vector<int> A(n * n);
    std::cout << "Original matrix A:\n";
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            A[i * n + j] = dist(gen);
            std::cout << A[i * n + j] << "\t";
        }
        std::cout << "\n";
    }

    std::vector<int> sorted_A;
    if (mode == "prolific") {
        prolific_scheduler scheduler(n); //n workers
        sorted_A = scheduler.sort_matrix_rows(A, n, n);
    } else if (mode == "collective") {
        collective_scheduler scheduler(n); // n workers
        sorted_A = scheduler.sort_matrix_rows(A, n, n);
    } else {
        std::cerr << "Invalid mode. Use 'prolific' or 'collective'.\n";
        return 1;
    }

    std::cout << "\nSorted matrix A (each row sorted):\n";
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            std::cout << sorted_A[i * n + j] << "\t";
        }
        std::cout << "\n";
    }

    return 0;
}