#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>

#define MASTER 0

void distribute_elements(int *vector, int *sub_vector, int n, int sub_n, int world_rank, int world_size) {
    // Distribute parts of the vector to all processes
    MPI_Scatter(vector, sub_n, MPI_INT, sub_vector, sub_n, MPI_INT, MASTER, MPI_COMM_WORLD);
}

void calculate_ranks(int *sub_vector, int *vector, int *ranks, int n, int sub_n, int world_rank) {
    // Calculate ranks for each element in the sub-vector
    for (int i = 0; i < sub_n; i++) {
        ranks[i] = 0;
        for (int j = 0; j < n; j++) {
            if (sub_vector[i] > vector[j] || (sub_vector[i] == vector[j] && (world_rank * sub_n + i) > j)) {
                ranks[i]++;
            }
        }
    }
}

void gather_and_sort(int *sorted_vector, int *sub_vector, int *ranks, int sub_n, int world_rank, int world_size) {
    // Gather the ranks from all processes
    int *all_ranks = NULL;
    if (world_rank == MASTER) {
        all_ranks = (int*) malloc(world_size * sub_n * sizeof(int));
    }
    MPI_Gather(ranks, sub_n, MPI_INT, all_ranks, sub_n, MPI_INT, MASTER, MPI_COMM_WORLD);

    // Master process sorts the array based on the ranks
    if (world_rank == MASTER) {
        for (int i = 0; i < world_size * sub_n; i++) {
            sorted_vector[all_ranks[i]] = sub_vector[i];
        }
    }

    free(all_ranks);
}

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);

    int world_rank, world_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    int n = 10; // Size of the vector
    int *vector, *sorted_vector, *sub_vector, *ranks;
    int sub_n = n / world_size; // Size of the sub-vector for each process

    if (world_rank == MASTER) {
        vector = (int*) malloc(n * sizeof(int));
        sorted_vector = (int*) malloc(n * sizeof(int));
        // Initialize the vector with random numbers
        for (int i = 0; i < n; i++) {
            vector[i] = rand() % 100; // Random numbers between 0 and 99
        }
    }

    sub_vector = (int*) malloc(sub_n * sizeof(int));
    ranks = (int*) malloc(sub_n * sizeof(int));

    distribute_elements(vector, sub_vector, n, sub_n, world_rank, world_size);
    calculate_ranks(sub_vector, vector, ranks, n, sub_n, world_rank);
    gather_and_sort(sorted_vector, sub_vector, ranks, sub_n, world_rank, world_size);

    if (world_rank == MASTER) {
        printf("Sorted array: ");
        for (int i = 0; i < n; i++) {
            printf("%d ", sorted_vector[i]);
        }
        printf("\n");
        free(vector);
        free(sorted_vector);
    }

    free(sub_vector);
    free(ranks);

    MPI_Finalize();
    return 0;
}