// Compile with: gcc -o gmp_miller-rabin_strong-lucas gmp_miller-rabin_strong_lucas.c -lm -lgmp
// Does A times Miller-Rabin PRP Tests followed by B times strong Lucas PRP Tests

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

int get_d_and_quotient(mpz_t quotient, mpz_t n) {
    int d=0;
    mpz_sub_ui(quotient, n, 1);
    while (mpz_even_p(quotient)) {
        mpz_tdiv_q_2exp(quotient, quotient, 1);
        d++;
    }
    return(d);
}

int get_e_and_quotient(mpz_t quotient, mpz_t n_plus_1) {
    int e=0;
    mpz_set(quotient, n_plus_1);
    while (mpz_even_p(quotient)) {
        mpz_tdiv_q_2exp(quotient, quotient, 1);
        e++;
    }
    return(e);
}

int miller_rabin(mpz_t base, mpz_t n, mpz_t align, int binary_length, int d, mpz_t quotient) {
    int i;
    mpz_t n_minus_1;
    mpz_t res;
    mpz_init(n_minus_1);
    mpz_init_set(res, base);
    mpz_sub_ui(n_minus_1, n, 1);
    for (i=binary_length-1; i>d; i--) {
        mpz_mul(res, res, res);
        mpz_mod(res, res, align);
        if (mpz_tstbit(n, i)) mpz_mul( res, res, base);
    }
    mpz_mul(res, res, res);
    if (mpz_tstbit(n, d)) mpz_mul( res, res, base);
    mpz_mod(res, res, n);
    if (mpz_cmp_ui(res, 1)==0 || mpz_cmp(res, n_minus_1)==0) {
        mpz_clear(res);
        mpz_clear(n_minus_1);
        return(1);
    }
    for (i=d-1; i>0; i--) {
        mpz_mul(res, res, res);
        mpz_mod(res, res, n);
        if (mpz_cmp(res, n_minus_1)==0) {
            mpz_clear(res);
            mpz_clear(n_minus_1);
            return(1);
        }
    }
    mpz_clear(res);
    mpz_clear(n_minus_1);
    return(0);
}

int strong_lucas(mpz_t a, mpz_t n, mpz_t align, mpz_t n_plus_1, int binary_length, int e, mpz_t quotient) {
    #define CLEAN_UP mpz_clear(u);mpz_clear(v);mpz_clear(A);mpz_clear(neg_A);mpz_clear(two);        mpz_clear(neg_2);

    int i;
    mpz_t u;
    mpz_t v;
    mpz_t A;
    mpz_t neg_A;
    mpz_t neg_2;
    mpz_t two;

    mpz_init_set_ui(u, 2);
    mpz_init_set(v, a);
    mpz_init_set(A, a);
    mpz_init(neg_A);
    mpz_sub(neg_A, n, a);
    mpz_init(neg_2);
    mpz_sub_ui(neg_2, n, 2);
    mpz_init_set_ui(two, 2);

    for (i=binary_length; i>=e; i--) {
        if (mpz_tstbit(n_plus_1, i)) {
            mpz_mul(u, u, v);
            mpz_sub(u, u, A);
            mpz_mul(v, v, v);
            mpz_sub_ui(v, v, 2);
        } else {
            mpz_mul(v, u, v);
            mpz_sub(v, v, A);
            mpz_mul(u, u, u);
            mpz_sub_ui(u, u, 2);
        }
        mpz_mod(u, u, align);
        mpz_mod(v, v, align);
    }
    if (mpz_congruent_p(u, two, n) && mpz_congruent_p(v, A, n)) {
        CLEAN_UP
        return(1);
    }
    for (i=e; i>0; i--) {
        if (mpz_congruent_p(u, neg_2, n) && mpz_congruent_p(v, neg_A, n)) {
            CLEAN_UP
            return(1);
        }
        mpz_mul(v, u, v);
        mpz_sub(v, v, A);
        mpz_mul(u, u, u);
        mpz_sub_ui(u, u, 2);
        mpz_mod(u, u, n);
        mpz_mod(v, v, n);
    }
    CLEAN_UP
    return(0);
}

int main(int argc, char* argv[]) {
    char s[1000000], *string=s;
    int i, binary_length, d, e; // d is the number of trailing zeroes of n-1; e of n+1
    int number_of_miller_rabin_tests;
    int number_of_strong_lucas_tests;
    mpz_t a;
    mpz_t n;
    mpz_t align;
    mpz_t gcd;
    mpz_t quotient;
    mpz_t discriminant;
    mpz_t n_plus_1;
 
    if (argc!=3) {
        printf("Usage:                      echo n | ./gmp_miller-rabin_strong-lucas number_of_fermat_tests number_of_lucas_tests\n");
        printf("Or like:   echo \"print(n)\" | gp -q | ./gmp_miller-rabin_strong-lucas number_of_fermat_tests number_of_strong-lucas_tests\n");
        printf("Or like: ./pfgw64 -k -f0 -od -q\"n\" | ./gmp_miller-rabin_strong-lucas number_of_fermat_tests number_of_strong-lucas_tests\n");
        exit(-1);
    }

    gets(s);
    if ((string=strchr(s,':')) && string[1]==' ') string+=2;
    else string=s;

    mpz_init_set_str(n, s, 10);

    if ((mpz_cmp_ui( n, 2 ))==0) {
        printf("Is prime.");
        exit(0);
    }
    if (!mpz_mod_ui(a, n, 2)) {
         printf("Is even.\n");
         exit(0);
    }
    if (mpz_perfect_square_p(n)) {
        printf("Is a square number.\n");
        exit(0);
    }

    mpz_init(align);
    mpz_mul_2exp(align, n, 32-mpz_mod_ui(align, n, 32));

    number_of_miller_rabin_tests=atoi(argv[1]);
    mpz_init(gcd);
    mpz_init_set_ui(a, 2);
    mpz_init(quotient);
    d=get_d_and_quotient(quotient, n);
    binary_length=mpz_sizeinbase(n, 2)-1;
    for (i=0; i<number_of_miller_rabin_tests; i++) {
        mpz_gcd(gcd, a, n );
        if(mpz_cmp_ui(gcd, 1)==0) {
            if (miller_rabin(a, n, align, binary_length, d, quotient)) {
                printf("Passed Miller-Rabin PRP test with base ");
                mpz_out_str(stdout, 10, a);
                printf(".\n");
            } else {
                printf("Failed Miller-Rabin PRP test with base ");
                mpz_out_str(stdout, 10, a);
                printf(".\nTesting aborted.\n");
                exit(0);
            }
        }
        mpz_nextprime(a, a);
    }

    number_of_strong_lucas_tests=atoi(argv[2]);
    mpz_set_ui(a, 2);
    mpz_init(discriminant);
    mpz_init(n_plus_1);
    mpz_add_ui(n_plus_1, n, 1);
    e=get_e_and_quotient(quotient, n_plus_1);
    binary_length=mpz_sizeinbase(n_plus_1, 2)-1;
    for (i=0; i<number_of_strong_lucas_tests; i++) {
        mpz_add_ui(a, a, 1);
        mpz_mul(discriminant, a, a);
        mpz_sub_ui(discriminant, discriminant, 4);
        while (mpz_jacobi(discriminant, n)!=-1 ) {
            mpz_add_ui(a, a, 1);
            mpz_mul(discriminant, a, a);
            mpz_sub_ui(discriminant, discriminant, 4);
        }
        if (strong_lucas(a, n, align, n_plus_1, binary_length, e, quotient)) {
            printf("Passed strong Lucas PRP test over x^2-");
            mpz_out_str(stdout, 10, a);
            printf("*x+1.\n");
        } else {
            printf("Failed strong Lucas PRP test over x^2-");
            mpz_out_str(stdout, 10, a);
            printf("*x+1.\nTesting aborted.\n");
            exit(0);
        }
    }
    printf("Testing completed.\n");
}