#include <mpi.h>
#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>

void sieve_of_eratosthenes(int start, int end, std::vector<int>& primes) {
    std::vector<bool> is_prime(end - start + 1, true);
    int limit = std::sqrt(end);

    for (int i = 2; i <= limit; ++i) {
        int multiple_start = std::max(i * i, (start + i - 1) / i * i);
        for (int j = multiple_start; j <= end; j += i) {
            is_prime[j - start] = false;
        }
    }

    for (int i = start; i <= end; ++i) {
        if (i > 1 && is_prime[i - start]) {
            primes.push_back(i);
        }
    }
}

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

    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    if (argc < 2) {
        if (rank == 0) {
            std::cerr << "Usage: " << argv[0] << " <n>" << std::endl;
        }
        MPI_Finalize();
        return 1;
    }

    int n = std::atoi(argv[1]);
    double start_time = MPI_Wtime();

    // Distribute ranges among processes
    int local_start = (rank * n) / size + 1;
    int local_end = ((rank + 1) * n) / size;

    std::vector<int> primes;
    sieve_of_eratosthenes(local_start, local_end, primes);

    // Compute local maximum gap
    int local_max_gap = 0;
    for (size_t i = 1; i < primes.size(); ++i) {
        int gap = primes[i] - primes[i - 1];
        local_max_gap = std::max(local_max_gap, gap);
    }

    // Share boundary primes with neighboring processes
    int left_prime = (primes.empty() ? 0 : primes.front());
    int right_prime = (primes.empty() ? 0 : primes.back());
    int left_neighbor_prime = 0, right_neighbor_prime = 0;

    MPI_Request requests[4];
    MPI_Isend(&left_prime, 1, MPI_INT, rank - 1, 0, MPI_COMM_WORLD, &requests[0]);
    MPI_Isend(&right_prime, 1, MPI_INT, rank + 1, 1, MPI_COMM_WORLD, &requests[1]);
    MPI_Irecv(&left_neighbor_prime, 1, MPI_INT, rank - 1, 1, MPI_COMM_WORLD, &requests[2]);
    MPI_Irecv(&right_neighbor_prime, 1, MPI_INT, rank + 1, 0, MPI_COMM_WORLD, &requests[3]);

    MPI_Waitall(4, requests, MPI_STATUSES_IGNORE);

    // Adjust gaps with boundary primes
    if (left_neighbor_prime > 0 && !primes.empty()) {
        int boundary_gap = primes.front() - left_neighbor_prime;
        local_max_gap = std::max(local_max_gap, boundary_gap);
    }
    if (right_neighbor_prime > 0 && !primes.empty()) {
        int boundary_gap = right_neighbor_prime - primes.back();
        local_max_gap = std::max(local_max_gap, boundary_gap);
    }

    // Find global maximum gap using reduction
    int global_max_gap = 0;
    MPI_Reduce(&local_max_gap, &global_max_gap, 1, MPI_INT, MPI_MAX, 0, MPI_COMM_WORLD);

    double end_time = MPI_Wtime();
    if (rank == 0) {
        std::cout << "Max Gap: " << global_max_gap << std::endl;
        std::cout << "Time: " << end_time - start_time << " seconds" << std::endl;
    }

    MPI_Finalize();
    return 0;
}