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

int main(int argc, char** argv) {
    // Initialize the MPI environment
    MPI_Init(&argc, &argv);

    // Get the rank of the process
    int world_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

    // Get the total number of processes
    int world_size;
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    // The master process (rank 0) prepares the data
    const int elements_per_proc = 3;
    int *scatter_data = NULL;
    int *gather_data = NULL;
    if (world_rank == 0) {
        const int total_elements = elements_per_proc * world_size;
        scatter_data = (int *)malloc(total_elements * sizeof(int));
        gather_data = (int *)malloc(total_elements * sizeof(int));
        // Initialize the array with values
        for (int i = 0; i < total_elements; i++) {
            scatter_data[i] = i + 1;
        }
    }

    // Each process will receive a portion of the array from the master process
    int recv_data[elements_per_proc];

    // Scatter the data from the master process to all processes
    MPI_Scatter(scatter_data, elements_per_proc, MPI_INT, recv_data, elements_per_proc, MPI_INT, 0, MPI_COMM_WORLD);

    // Each process performs some operation on the received data
    for (int i = 0; i < elements_per_proc; i++) {
        recv_data[i] *= 2; // Example operation: multiply each element by 2
    }

    // Gather the modified data from all processes back to the master process
    MPI_Gather(recv_data, elements_per_proc, MPI_INT, gather_data, elements_per_proc, MPI_INT, 0, MPI_COMM_WORLD);

    // The master process prints the gathered data
    if (world_rank == 0) {
        printf("Gathered data: ");
        for (int i = 0; i < elements_per_proc * world_size; i++) {
            printf("%d ", gather_data[i]);
        }
        printf("\n");
    }

    // Clean up
    if (world_rank == 0) {
        free(scatter_data);
        free(gather_data);
    }

    // Finalize the MPI environment
    MPI_Finalize();
}