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

// Function to check if a number is prime
int is_prime(int num) {
    if (num <= 1) return 0; // 0 and 1 are not prime
    if (num <= 3) return 1; // 2 and 3 are prime
    if (num % 2 == 0 || num % 3 == 0) return 0; // eliminate even numbers and multiples of 3
    for (int i = 5; i <= sqrt(num); i += 6) {
        if (num % i == 0 || num % (i + 2) == 0) return 0;
    }
    return 1;
}

int main(int argc, char** argv) {
    int rank, size;
    int lower, upper;

    // Initialize MPI
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // Read the range of numbers from the user (only by the master process)
    if (rank == 0) {
        printf("Enter the lower bound of the range: ");
        scanf("%d", &lower);
        printf("Enter the upper bound of the range: ");
        scanf("%d", &upper);
    }

    // Broadcast the range to all processes
    MPI_Bcast(&lower, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&upper, 1, MPI_INT, 0, MPI_COMM_WORLD);

    // Calculate the range of numbers for each process
    int range = upper - lower + 1;
    int local_lower = lower + (rank * range / size);
    int local_upper = lower + ((rank + 1) * range / size) - 1;

    // Find primes in the assigned range
    int* local_primes = malloc((local_upper - local_lower + 1) * sizeof(int));
    int local_count = 0;

    for (int i = local_lower; i <= local_upper; i++) {
        if (is_prime(i)) {
            local_primes[local_count++] = i;
        }
    }

    // Gather prime counts from all processes
    int* recv_counts = NULL;
    if (rank == 0) {
        recv_counts = malloc(size * sizeof(int));
    }
    MPI_Gather(&local_count, 1, MPI_INT, recv_counts, 1, MPI_INT, 0, MPI_COMM_WORLD);

    // Gather all primes from all processes
    int* displacements = NULL;
    int total_primes = 0;
    if (rank == 0) {
        displacements = malloc(size * sizeof(int));
        displacements[0] = 0;
        for (int i = 1; i < size; i++) {
            displacements[i] = displacements[i - 1] + recv_counts[i - 1];
        }
        total_primes = displacements[size - 1] + recv_counts[size - 1];
    }

    int* global_primes = NULL;
    if (rank == 0) {
        global_primes = malloc(total_primes * sizeof(int));
    }
    MPI_Gatherv(local_primes, local_count, MPI_INT, global_primes, recv_counts, displacements, MPI_INT, 0, MPI_COMM_WORLD);

    // Print the result
    if (rank == 0) {
        printf("Prime numbers in the range [%d, %d]:\n", lower, upper);
        for (int i = 0; i < total_primes; i++) {
            printf("%d ", global_primes[i]);
        }
        printf("\n");

        // Free dynamically allocated memory
        free(global_primes);
        free(recv_counts);
        free(displacements);
    }

    // Finalize MPI and free local memory
    free(local_primes);
    MPI_Finalize();

    return 0;
}