#include "nth_prime.h"

#include <algorithm>
#include <cmath>
#include <cstring>
#include <stdexcept>
#include <vector>


namespace prime {

using namespace std;

number_t nth(size_t n, size_t sieve_size)
{
    if (n < 1)
        throw domain_error("invalid prime number index");

    if (n == 1)
        return 2;

    vector<number_t> primes = {3, 5, 7, 11, 13, 17, 19, 23, 29, 31};

    if (n - 2 < primes.size())
        return primes[n - 2];

    primes.reserve(n - 1);

    if (auto_sieve_size == sieve_size) {

        sieve_size = min(size_t(32 * 1024),
                         size_t(0.5*(n * (log(n) + log(log(n)))) + 0.5));
    }

    vector<char>     sieve(sieve_size, false);
    vector<number_t> next(n - 1, 0);

    sieve[0] = true;

    number_t first_number = 1;

    auto create_next = [&first_number](number_t p) -> number_t {

        return (p*p - first_number)/2;
    };

    transform(primes.cbegin(), primes.cend(), next.begin(), create_next);

    for (auto p : primes)
        sieve[(p - first_number) / 2] = true;

    auto process_prime = [sieve_size,&sieve](const number_t &p, number_t &next) {

        number_t i = next;

        for (; i < sieve_size; i += p) {

            sieve[i] = true;
        }

        next = i - sieve_size;
    };


    while (true) {

        for (size_t pn = 0; pn < primes.size(); ++pn) {

            const number_t &p = primes[pn];

            number_t i = next[pn];

            for (; i < sieve_size; i += p) {

                sieve[i] = true;
            }

            next[pn] = i - sieve_size;
//            process_prime(primes[pn], next[pn]);
        }

        for (size_t i = 0; i < sieve_size; ++i) {

            if (sieve[i])
                continue;

            number_t p = first_number + 2*i;
        
            primes.push_back(p);

            if (primes.size() == n - 1)
                return p;

            next[primes.size() - 1] = create_next(p);//(p*p - first_number)/2;

            sieve[i] = true;

            number_t j = next[primes.size() - 1];

            for (; j < sieve_size; j += p) {

                sieve[j] = true;
            }

            next[primes.size() - 1] = j - sieve_size;
        }

        number_t next_first_number = first_number + 2*sieve_size;

        if (next_first_number < first_number)
            throw runtime_error("failed to reach prime number");

        first_number = next_first_number;
        memset(&sieve[0], false, sieve.size()*sizeof(sieve[0]));
    }
}

}