#include <gmp.h>
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <time.h>


/* n is the number to be factored, and a is the random parameter in F */
mpz_t n, a;
unsigned long E; // The exponent in an eventual Fermat Number


/* The function to be iterated */
void F(mpz_t r, int f) {
    switch(f) {
        /* F(x) = (x² + a) mod n */
        case 1:
            mpz_powm_ui(r, r, 2, n);
            mpz_add(r, r, a);
            mpz_mod(r, r, n);
            break;
        
        /* F(x) = (x^E + 1) mod n, good for factoring Fermat Number */
        case 2:
            mpz_powm_ui(r, r, E, n);
            mpz_add_ui(r, r, 1);
            mpz_mod(r, r, n);
            break;
    }
}


/* Main method */
int main(int argc, char *argv[]) {
    gmp_randstate_t state;
    mpz_t s, high_bound, U, V, UV, g, p;
    mpz_init(a);
    mpz_init(s);
    mpz_init(high_bound);
    mpz_init(U);
    mpz_init(V);
    mpz_init(UV);
    mpz_init(g);
    unsigned long iterations=0;
    int f, i=1;
    
    /* Read the command line arguments */
    if(argc < 2) { // If we do not have any parameters
        printf("Using default input - the 8th Fermat number\n");
        mpz_ui_pow_ui(n, 2, 256);
        mpz_add_ui(n, n, 1);
        f = 2; // use a function F that is good at factoring Fermat numbers
        E = 256;
    } else {
        mpz_init_set_str(n,argv[argc-1],10);
        f = 1; // use the default function F

        while(i<argc) {
            /* The parameter '-F n' gives the n'th Fermat Number as input */
            if(!strcmp(argv[i], "-F") && i<argc-1) {
                E = pow(2, atoi(argv[i+1]));
                printf("Factoring the %dth Fermat Number.\n", atoi(argv[i+1]));
                mpz_ui_pow_ui(n, 2, E);
                mpz_add_ui(n, n, 1);
                f=2;
                i++;
            }                    

            /* Perform a primality test on n before trying to factor */
            if(!strcmp(argv[i],"-t")) {
                if(mpz_probab_prime_p(n, 5)) {
                    printf("Probably a prime...\n");
                    return 0;
                }
            }

            /* Display help */
            if(!strcmp(argv[i], "-h")) {
                printf("Usage: ./rho [options] [n]\n n     The number to be factored\n\nOptions:\n");
                printf(" -h    This help\n -F k  Use the k'th Fermat Number as input\n -t    Use Rabin-Miller to test for primality\n");
                return 0;
            }
            i++;
        }
    }

    gmp_printf("n=%Zd\n", n);

    /* Asymptotic running time is O(sqrt(p)), and p~sqrt(n) in worst case */
    mpz_sqrt(p, n);
    mpz_sqrt(p, p);
    gmp_printf("This should be done in ~%Zd iterations (in absolute worst case).\n", p);

     /* Give the randomizer the timer as a seed */
     srand( (unsigned int) time( NULL ));
    gmp_randinit_default(state);
    gmp_randseed_ui(state, random());
    mpz_sub_ui(high_bound, n, 3);

    do {
        /* Pick random a in [1,n-3] and s in [0,n-1] */
        mpz_urandomm(s, state, n);
        if(f=1) {
            mpz_urandomm(a, state, high_bound);
            mpz_add_ui(a, a, 1);
        }
        /* U=V=s */
        mpz_set(U, s);
        mpz_set(V, s);

        do {
            /* U = F(U) and V = F(F(V)). g = gcd(U-V, n) */
            F(U,f);
            F(V,f);
            F(V,f);
            mpz_sub(UV, U, V);
            mpz_gcd(g, UV, n);
            
            /* Number of iterations done so far */
            iterations++;
    
            /* Print a dot for every 100000 iteration */        
            if(iterations%100000 == 0) {
                printf(".");
                fflush(stdout);
            }
        } while(mpz_cmp_ui(g,1) == 0);
    } while(mpz_cmp(g,n) == 0);

    gmp_printf("\nNon-trivial divisor found!\nd=%Zd\n", g);
    printf("Iterations: %d\n", iterations);
}