// 1. 变量声明放外部
// 2. 减少乘除法为加法
// 3. 乘或除以2改为移位方式
#include "mpi.h"
#include <math.h>
#include <stdio.h>
#include <cstring>
#define MIN(a,b)  ((a)<(b)?(a):(b))
#define VALUE_TO_INDEX(x) (x - 3 >> 1)
#define INDEX_TO_VALUE(x) ((x << 1) + 3)

#define block_size 589824   // 37748736 / 64

typedef long long LL;

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;
    LL     j;
    int    id;           /* Process ID number */
    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' */
    LL     low_index;
    LL     high_index;
    char  *st;           /* Get the prime */
    int   *primes;
    LL     block_low_index;
    LL     block_high_index;
    LL     block_low_value;
    LL     block_high_value;
    LL     prime_pos;
    LL     prime_pos_square;
    LL     t;

    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_index = 0 + ((LL) id * (n - 1) / p >> 1);            // index为3开始的奇数映射的下标，3-->0，5-->1，依此类推
    high_index = -1 + ((LL) (id + 1) * (n - 1) / p >> 1);
    low_value = INDEX_TO_VALUE(low_index);
    high_value = INDEX_TO_VALUE(high_index);
    size = high_index - low_index + 1;

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

    marked = (char *) malloc (size + 10);

    if (marked == NULL) {
        printf ("Cannot allocate enough memory for marked\n");
        MPI_Finalize();
        exit (1);
    }

    // for (i = 0; i < size; i++) marked[i] = 0;
    memset(marked, 0, size + 10);

    LL N = (LL) sqrt(n);
    st = (char *) malloc (N + 10);
    primes = (int *) malloc (N + 10);
    int primes_cnt = 0;

    if (st == NULL || primes == NULL) {
        printf ("Cannot allocate enough memory for st or primes\n");
        MPI_Finalize();
        exit (1);
    }

    // for (i = 0; i < N + 10; i++) st[i] = 0;
    memset(st, 0, N + 10);

    // for (i = 2; i <= N; i++) {
    //     if (!st[i]) primes[primes_cnt++] = i;
    //     for (j = 0; primes[j] <= N / i; j++) {
    //         st[primes[j] * i] = 1;
    //         if (i % primes[j] == 0) break;
    //     }
    // }
    
    for (i = 2; i <= N; i++) {
        if (st[i]) continue;
        primes[primes_cnt++] = i;
        for (j = i * i; j <= N; j += i) st[j] = 1;
    }

    block_low_index = low_index;
    block_high_index = block_size - low_index % block_size + low_index;
    block_low_value = INDEX_TO_VALUE(block_low_index);
    block_high_value = INDEX_TO_VALUE(block_high_index);
    
    while (block_low_index <= high_index) {
        if (block_high_index > high_index) {
            block_high_index = high_index;
            block_high_value = high_value;
        }

        for (i = 1; i < primes_cnt && primes[i] * primes[i] <= high_value; i++) {
            prime_pos = primes[i];
            prime_pos_square = prime_pos * prime_pos;

            if (prime_pos_square > block_high_value) continue;

            if (prime_pos_square >= block_low_value)
                first = VALUE_TO_INDEX(prime_pos_square);
            else {
                if (!(block_low_value % prime_pos)) first = block_low_index;
                else {
                    t = prime_pos - (block_low_value % prime_pos) + block_low_value;
                    if (t % 2) first = VALUE_TO_INDEX(t);
                    else first = VALUE_TO_INDEX(t + prime_pos);
                }
            }

            for (j = first; j <= MIN(block_high_index, block_low_index + size - 1); j += prime_pos) {
                marked[j - low_index] = 1;
            }
        }

        block_low_index = block_high_index + 1;
        block_low_value = block_high_value + 2;
        block_high_index += block_size;
        block_high_value += block_size << 1;
    }

    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,
        0, MPI_COMM_WORLD);
    else global_count = count;

    /* Stop the timer */

    elapsed_time += MPI_Wtime();

    /* Print the results */
    if (n >= 2) global_count++;
    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;
}