#include <stdexcept>
#include <cstring>
#include <cmath>
#include <cassert>
#include "numalgo.h"
#include "primes.h"

using namespace std;

vector<int> smallprimes;
bool smallprimes_inited=false;

void fillsmallprimes(){
	// cerr<<"Filling small primes..."<<endl;
	smallprimes_inited=true;
	//TODO: reserve expected amount of space in smallprimes
	smallprimes.push_back(2);
	const int highbound=65000;
	bool composite[highbound/2]; //entries for 3, 5, 7, 9, etc.
	memset(composite,0,highbound/2*sizeof(bool));
	int roothighbound=sqrt(highbound);
	for(int i=3;i<=highbound;i+=2){
		if(composite[i/2-1])continue;
		smallprimes.push_back(i);
		if(i>roothighbound)continue;
		for(int j=i*i;j<=highbound;j+=2*i){
			composite[j/2-1]=true;
		}
	}
	//for(int p : smallprimes)cerr<<p<<' ';
	//cerr<<endl;
}

pair<Bigint,Bigint> genprimepair(int nbits){
	// for x = nbits/2:
	// (2^x)^2 = 2^(2x)
	// (2^x + 2^(x-2))^2 = 2^(2x) + 2^(2x-1) + 2^(2x-4)
	// ergo: (2^x + lambda*2^(x-2))^2 \in [2^(2x), 2^(2x+1)), for lambda \in [0,1]
	// To make sure the primes "differ in length by a few digits" [RSA78], we use x1=x-2 in the first
	// prime and x2-x+2 in the second
	int x1=nbits/2-2,x2=(nbits+1)/2+2;
	assert(x1+x2==nbits);
	return make_pair(
		randprime(Bigint::one<<x1,(Bigint::one<<x1)+(Bigint::one<<(x1-2))),
		randprime(Bigint::one<<x2,(Bigint::one<<x2)+(Bigint::one<<(x2-2))));
}

Bigint randprime(const Bigint &biglow,const Bigint &bighigh){
	//https://en.wikipedia.org/wiki/Generating_primes#Large_primes
	if(!smallprimes_inited)fillsmallprimes();

	assert(bighigh>biglow);
	static const int maxrangesize=100001;
	Bigint diff(bighigh-biglow);
	Bigint low,high; //inclusive
	if(diff<=maxrangesize){
		low=biglow;
		high=bighigh;
		// cerr<<"low=biglow="<<low<<" high=bighigh="<<high<<endl;
	} else {
		high=low=cryptrandom_big(diff-maxrangesize);
		high+=maxrangesize;
		// cerr<<"low="<<low<<" high="<<high<<endl;
	}
	if(low.even())low+=1;
	if(high.even())high-=1;
	// cerr<<"low="<<low<<" high="<<high<<endl;
	Bigint sizeb(high-low+1);
	assert(sizeb>=0&&sizeb<=maxrangesize);
	int nnums=(sizeb.lowdigits()+1)/2;
	// cerr<<"nnums="<<nnums<<endl;

	int nleft=nnums;

	bool composite[nnums]; //low, low+2, low+4, ..., high (all odd numbers)
	memset(composite,0,nnums*sizeof(bool));
	int nsmallprimes=smallprimes.size();
	for(int i=1;i<nsmallprimes;i++){
		int pr=smallprimes[i];
		if(pr>=2*nnums)break;
		int lowoffset=low.divmod(Bigint(pr)).second.lowdigits();
		int startat;
		if(lowoffset==0)startat=0;
		else if((pr-lowoffset)%2==0)startat=(pr-lowoffset)/2;
		else startat=(pr-lowoffset+pr)/2;
		// cerr<<"pr="<<pr<<" lowoffset="<<lowoffset<<" startat="<<startat<<endl;
		for(int i=startat;i<nnums;i+=pr){ //skips ahead `2*pr` each time (so `pr` array elements)
			// cerr<<"sieving composite["<<i<<"]: "<<low+2*i<<endl;
			nleft-=!composite[i];
			composite[i]=true;
		}
	}
	vector<int> maybeprimes;
	maybeprimes.reserve(nleft);
	// cerr<<"Left ("<<nleft<<"):";
	for(int i=0;i<nnums;i++){
		if(!composite[i]){
			// cerr<<' '<<low+2*i;
			maybeprimes.push_back(i);
		}
	}
	// cerr<<endl;

	while(maybeprimes.size()){
		int idx=arc4random_uniform(maybeprimes.size());
		int i=maybeprimes[idx];
		Bigint bi(low+2*i);
		if(bailliePSW(bi))return bi;
		maybeprimes.erase(maybeprimes.begin()+idx);
	}
	throw range_error("No primes");
}

bool strongPseudoPrime2(const Bigint &n){
	//https://en.wikipedia.org/wiki/Strong_pseudoprime#Formal_definition
	if(n<2||n.even())return false;
	Bigint nm1(n);
	nm1-=1;
	Bigint d(nm1);
	while(d.even()){ //TODO: optimise using __builtin_ctz
		d>>=1;
		if(expmod(Bigint::two,d,n)==nm1)return true;
	}
	if(expmod(Bigint::two,d,n)==1)return true;
	return false;
}

bool strongLucasPrime(const Bigint &n){
	//https://en.wikipedia.org/wiki/Lucas_pseudoprime#Implementing_a_Lucas_probable_prime_test
	//TODO: use d and s already calculated in strongPseudoPrime2
	if(n.even())return false;
	int D;
	if     ((D=  5),jacobiSymbol(Bigint(D),n)==-1);
	else if((D= -7),jacobiSymbol(Bigint(D),n)==-1);
	else if((D=  9),jacobiSymbol(Bigint(D),n)==-1);
	else if((D=-11),jacobiSymbol(Bigint(D),n)==-1);
	else {
		Bigint root(isqrt(n));
		if(root*root==n)return false; //perfect square
		for(int i=6;;i++){
			D=2*i+1;
			if(i%2==1)D=-D;
			if(jacobiSymbol(Bigint(D),n)==-1)break;
		}
	}
	//now we have a D
	// cerr<<"D="<<D<<endl;
	int P=1,iQ=(1-D)/4;
	if(gcd(n,Bigint(P))!=1||gcd(n,Bigint(iQ))!=1)return false; //would be too easy to ignore

	//now begin the actual sequence algorithm
#if 1
	int s=0;
	Bigint d(n);
	d-=1;
	while(d.even()){
		d>>=1;
		s++;
	}
#else
	Bigint d(n);
	d+=1;
#endif
	// cerr<<"d="<<d<<endl;
	vector<bool> dbits=d.bits();
	assert(dbits.size()>0);
	assert(dbits[dbits.size()-1]==true);
	Bigint U(1),V(P);
	Bigint Qk(iQ);
	for(int i=dbits.size()-2;i>=0;i--){
		U*=V;
		V*=V;
		V-=Qk<<1;
		Qk*=Qk;
		if(dbits[i]){
			Bigint unext(U*P);
			unext+=V;
			if(unext.odd())unext+=n;
			assert(unext.even());
			unext>>=1;
			Bigint vnext(U*D);
			vnext+=V*P;
			if(vnext.odd())vnext+=n;
			assert(vnext.even());
			vnext>>=1;
			U=unext.divmod(n).second;
			V=vnext.divmod(n).second;
			Qk=(Qk*iQ).divmod(n).second;
		} else {
			U=U.divmod(n).second;
			V=V.divmod(n).second;
			Qk=Qk.divmod(n).second;
		}
	}
	if(U==0)return true;
	if(V==0)return true; //r=0 check
#if 1
	for(int r=1;r<s;r++){
		V*=V;
		V-=Qk<<1;
		V=V.divmod(n).second;
		Qk=(Qk*Qk).divmod(n).second;
		if(V==0)return true;
	}
#endif
	return false;
}

bool bailliePSW(const Bigint &n){
	//https://en.wikipedia.org/wiki/Baillie%E2%80%93PSW_primality_test#The_test
	return strongPseudoPrime2(n)&&strongLucasPrime(n);
}