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

int test_n_against_a(mpz_t n, mpz_t n_1, mpz_t d, mp_bitcnt_t s, mpz_t a, mpz_t x);


int main(int argc, char *argv[]) {
    int world_size;
    int world_rank;
    mpz_t n;
    mpz_t n_1;
    mp_bitcnt_t s;
    mpz_t d;
    mpz_t a;
    mpz_t a_max;
    mpz_t x;
    size_t logn;
    int answer_local; // 0: inconclusive, 1: prime, 2: composite
    int answer_global; // 0: inconclusive, 1: prime, 2: composite

    MPI_Request answer_found_req;
    int answer_found_flag;
    int is_composite;


    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

    // print an error if the starting position is not specified
    if (argc != 2) {
        printf("USAGE: %s <start>\n", argv[0]);
        printf("  start: an integer as a strating point (arbitrary precision)\n", argv[0]);
        MPI_Finalize();
        exit(EXIT_FAILURE);
    }


    // print an error if working with a single process
    if (world_size < 2) {
        printf("ERROR: please create at least 2 processes!\n");
        MPI_Finalize();
        exit(EXIT_FAILURE);
    }


    mpz_init_set_str(n, argv[1], 10);

    mpz_init(n_1);
    mpz_init(d);
    mpz_init(a);
    mpz_init(a_max);
    mpz_init(x);


    while(1) {
        // PREPARE VARIABLES

        // n_1 = n - 1
        mpz_sub_ui(n_1, n, 1);

        // s = [ x WHERE n-1 = 2^x * d ]
        s = mpz_scan1(n_1, 0);

        // d = n_1 / 2^s
        mpz_tdiv_q_2exp(d, n_1, s);

        // logn = log(n)
        logn = mpz_sizeinbase(n, 2);

        // a_max = ( (logn^2) / [number of workers] )
        mpz_set_ui(a_max, logn);
        mpz_mul_ui(a_max, a_max, logn);
        // TODO: range starts from 2, we scan a bit extra
        mpz_cdiv_q_ui(a_max, a_max, world_size);

        // a = [start of scanning range for this worker]
        // a_max = [end of scanning range for this worker]
        mpz_mul_ui(a, a_max, world_rank);
        mpz_mul_ui(a_max, a_max, world_rank + 1);
        // range should start from 2
        mpz_add_ui(a, a, 3);
        mpz_add_ui(a_max, a_max, 3);

        answer_local = 0;

#ifdef DEBUG
        if (world_rank == 0) {
            gmp_printf("n     : %Zd\n", n);
            gmp_printf("n-1   : %Zd\n", n_1);
            gmp_printf("s     : %lu\n", s);
            gmp_printf("d     : %Zd\n", d);
            gmp_printf("logn  : %d\n", logn);
        }
        gmp_printf("p%02d   : %Zd - %Zd\n", world_rank, a, a_max);
#endif

        // scan the assigned range
        do {

            // can it be proven that n is composite?
            if (test_n_against_a(n, n_1, d, s, a, x)) {
                // n is composite
                answer_local = 2;
            } 

            // is {the local answer still inconclusive} and {test range is exhausted}?
            if ( (answer_local == 0) && ((mpz_cmp(a, a_max) >= 0))) {
                // n might be a prime
                answer_local = 1;
            }

#ifdef DEBUG
            printf("composition: %d found %d\n", world_rank, answer_local);
#endif

            // reduce local answers. more definite answers override others
            // {2: definitely composite} > {1: might be prime} > {0: inconclusive}
            MPI_Allreduce(&answer_local, &answer_global, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD);

            // prepare to check against the next a in range
            mpz_add_ui(a, a, 1);

            // answer_global == {0 :inconclusive} : inconclusive, restart to check against the next a in range
            // answer_global == {1: might be prime} : inconclusive, but no more iterations available, exit
            // answer_global == {2: definitely composite} : conclusive, exit
        } while (answer_global == 0); 

        // print out the agreed answer (only on the first worker)
        if (world_rank == 0) {
#ifdef DEBUG
            gmp_printf(" %d : %Zd\n", answer_global, n);
#else
            if (answer_global == 1) {
                gmp_printf("%Zd\n", n);
            }
#endif
        }

        // increment n
        mpz_add_ui(n, n, 1);

    }

    MPI_Finalize();
}




int test_n_against_a(mpz_t n, mpz_t n_1, mpz_t d, mp_bitcnt_t s, mpz_t a, mpz_t x) {
    // 0: inconclusive
    // 1: number is composite

    mp_bitcnt_t r;

    // x = a^d (mod n)
    mpz_powm(x, a, d, n);

    // if (x == 1) return inconclusive
    if ( ! mpz_cmp_ui(x, 1) ) return 0;

    // if (x == n-1) return inconclusive
    if ( ! mpz_cmp(x, n_1) ) return 0;

    for (r=1; r<s; r++) {

        // x = x^2 (mod n)
        mpz_powm_ui(x, x, 2, n);

        // if (x == n-1) return inconclusive
        if ( ! mpz_cmp(x, n_1) ) return 0;

    }

    // return is_composite
    return 1;

}