#include "mpi.h"
#include <math.h>
#include <stdio.h>
#define MIN(a,b)  ((a)<(b)?(a):(b))
typedef long long LL;   // 防止入参n在1e10的情况下爆int

int main (int argc, char *argv[])
{
    LL     count;        /* Local prime count */
    double elapsed_time; /* Parallel execution time */
    LL     first;        /* Index of first multiple */
    LL     global_count; /* Global prime count */
    LL     high_value;   /* Highest value on this proc */
    LL     i;
    int    id;           /* Process ID number */
    LL     index;        /* Index of current prime */
    LL     low_value;    /* Lowest value on this proc */
    char  *marked;       /* Portion of 2,...,'n' */
    LL     n;            /* Sieving from 2, ..., 'n' */
    int    p;            /* Number of processes */
    int    proc0_size;   /* Size of proc 0's subarray */
    LL     prime;        /* Current prime */
    LL     size;         /* Elements in 'marked' */

    MPI_Init (&argc, &argv);

    /* Start the timer */

    MPI_Comm_rank (MPI_COMM_WORLD, &id);
    MPI_Comm_size (MPI_COMM_WORLD, &p);
    MPI_Barrier(MPI_COMM_WORLD);
    elapsed_time = -MPI_Wtime();

    if (argc != 2) {
        if (!id) printf ("Command line: %s <m>\n", argv[0]);
        MPI_Finalize();
        exit (1);
    }

    n = atoll(argv[1]);

    /* Figure out this process's share of the array, as
        well as the integers represented by the first and
        last array elements */

    low_value = 2 + (LL) id * (n - 1) / p;              // 使用块分配策略，每个进程得到自己负责的value上下界
    high_value = 1 + (LL) (id + 1) * (n - 1) / p;
    size = high_value - low_value + 1;                  // size为负责的范围大小

    /* Bail out if all the primes used for sieving are
        not all held by process 0 */

    proc0_size = (n - 1)/p;

    if ((2 + proc0_size) < (int) sqrt((double) n)) {    // 这里是检测需要的素数是否都在0号进程的范围之内，若否则报错
        if (!id) printf ("Too many processes\n");
        MPI_Finalize();
        exit (1);
    }

    /* Allocate this process's share of the array. */

    marked = (char *) malloc (size);                    // 每个进程申请自己所需大小的空间

    if (marked == NULL) {                               // 若无法申请到足够的空间，则报错
        printf ("Cannot allocate enough memory\n");
        MPI_Finalize();
        exit (1);
    }

    for (i = 0; i < size; i++) marked[i] = 0;           // 将标记数组初始化
    if (!id) index = 0;                                 // index == 0 对应着：当前的prime == 2
    prime = 2;                                          // 2为第一个素数
    do {
        // first为每一个进程所负责的范围内的，第一个为当前素数的倍数
        if (prime * prime > low_value)
            first = prime * prime - low_value;
        else {
            if (!(low_value % prime)) first = 0;
            else first = prime - (low_value % prime);
        }
        for (i = first; i < size; i += prime) marked[i] = 1;    // 将当前素数的所有倍数都筛掉
        if (!id) {
            while (marked[++index]);                            // 找到下一个素数
            prime = index + 2;
        }
        if (p > 1) MPI_Bcast (&prime,  1, MPI_LONG_LONG_INT, 0, MPI_COMM_WORLD);    // 广播该素数
    } while (prime * prime <= n);
    count = 0;
    for (i = 0; i < size; i++)          // 统计所有进程中未被筛掉的数（素数）的个数
        if (!marked[i]) count++;
    if (p > 1) MPI_Reduce (&count, &global_count, 1, MPI_LONG_LONG_INT, MPI_SUM,    // 将其他进程的count归约到0号进程
        0, MPI_COMM_WORLD);
    else global_count = count;

    /* Stop the timer */

    elapsed_time += MPI_Wtime();


    /* Print the results */

    if (!id) {
        printf ("There are %lld primes less than or equal to %lld\n",
            global_count, n);
        printf ("SIEVE (%d) %10.6f\n", p, elapsed_time);
    }
    MPI_Finalize ();
    return 0;
}