// The routine roots64s computes the roots using our Tangent Graeffe
// root finding C implementation.
// Copyright Michael Monagan March 2020

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

#define DEBUG 0
#define LONG long long int
#define ULONG unsigned long long int
unsigned LONG seed;
unsigned LONG mult;
LONG rand64s(LONG p);
void randvec64s( LONG *A, int n, LONG p );
LONG * array( int n );
LONG * array64s( LONG n );

#define UINT64 unsigned long long

typedef struct {
  UINT64 s;   /* shift */
  UINT64 v;   /* reciprocal of d */
  UINT64 d0;  /* divisor shifted up */
  UINT64 d1;
} recint;
recint recip1(UINT64  p);
UINT64 mulrec64(UINT64  a, UINT64  b, recint  v);


/******************************************************************************************/
/*  Zp utilities                                                                          */
/******************************************************************************************/

#define max32s max32smain
#define min32s min32smain
#define max64s max64smain
#define min64s min64smain
#define add64s add64smain
#define sub64s sub64smain
#define mul64s mul64smain
#define neg64s neg64smain
LONG min64s( LONG a, LONG b ) { if( a<b ) return a; else return b; }
LONG max64s( LONG a, LONG b ) { if( a>b ) return a; else return b; }
LONG neg64s(LONG a, LONG p) { return (a==0) ? 0 : p-a; }
LONG add64s(LONG a, LONG b, LONG p) { LONG t; t = (a-p)+b; t += (t>>63) & p; return t; }
LONG sub64s(LONG a, LONG b, LONG p) { LONG t; t = a-b; t += (t>>63) & p; return t; }
LONG mul64s(LONG a, LONG b, LONG p) {
        LONG q, r;
        __asm__ __volatile__(           \
        "       mulq    %%rdx           \n\t" \
        "       divq    %4              \n\t" \
        : "=a"(q), "=d"(r) : "0"(a), "1"(b), "rm"(p));
        return r;
}

LONG modinv64s( LONG c, LONG p );
LONG powmod64s( LONG a, LONG n, LONG p );

/******************************************************************************************/
/*  Polynomial utitlities from polyalg2                                                   */
/******************************************************************************************/

LONG poleval64s(LONG *a, int d, LONG x, LONG p);
int polmul64s( LONG * A, LONG * B, LONG * C, int da, int db, LONG p);
int poldiv64s( LONG * A, LONG * B, int da, int db, LONG p );
void vecprint64s( LONG *A, int n );
int vecequal64s( LONG *A, LONG *B, int n );
void polprint64s( LONG *A, int d, LONG p );
int polequal64s( LONG *a, LONG *b, int d );
void polcopy64s( LONG *A, int d, LONG *B );
int poldiff64s( LONG *f, int d, LONG *fp, LONG p );
void polcopy64s( LONG *A, int d, LONG *B );
void FFTgraeffe64s( LONG * f, LONG * g, LONG d, LONG *T, LONG p );
LONG getomega64s( LONG p, LONG n );
LONG poleval64s(LONG *a, int d, LONG x, LONG p);
LONG bluestein64s( LONG *a, LONG n, LONG s, LONG *v, LONG p );
LONG FFTbluestein64s( LONG *a, LONG n, LONG s, LONG *v, LONG *T, LONG p );
void tangGraeffe64s( LONG *f, LONG *g, LONG d, LONG p );
void FFTtangGraeffe64s( LONG *f, LONG *g, LONG d, int k, LONG *T, LONG p );
int FFTpoldiv64s( LONG *a, LONG *b, LONG da, LONG db, long p );
void fastLambda( LONG * v, LONG n, LONG *f, LONG p );
void changebase( LONG *f, LONG n, LONG alpha, LONG p );
void changebase2( LONG *f, LONG n, LONG alpha, LONG p );
void changebase3( LONG *f, LONG n, LONG alpha, LONG p );


void Lambda( LONG *v, LONG n, LONG *f, LONG p ) {
    LONG i;
    LONG a[2];
    a[1] = 1;
    f[0] = 1;
    for( i=0; i<n; i++ ) {
        a[0] = neg64s(v[i],p); // a = x-v[i]
        polmul64s(a,f,f,1,i,p);
    }
    return;
}


void translate( LONG *f, LONG n, LONG *g, LONG p ) {
// calculate f(x+e) mod e^2 = f(x) + g(x) e where deg(f) = n
// This is a waste of time since g(x) = f'(x)
   LONG i,m; LONG *f1, *f2, *g1, *g2;
   if( n==0 ) return; // g = 0
   // f = a+bx  g(x) = coeff(f(x+e),e) = b = f[1]
   if( n==1 ) { g[0] = f[1]; return; }
   // f = a+bx+cx^2  g(x) = coeff(f(x+e),e) = b+2cx = f[1]+2f[2]
   if( n==2 ) { g[0] = f[1]; g[1] = add64s(f[2],f[2],p); return; } //
   m = n/2+1;
   f1 = f; // degree m-1
   f2 = f+m; // degree n-m
   for( i=0; i<n; i++ ) g[i] = 0;
   g1 = g; // degree m-2
   translate(f1, m-1, g1, p );
   //g2 = T; // degree (n-m-1)
   g2 = g+m;
   // n=6, m=4, deg(f1)=3, deg(f2)=3, deg(g1)=2, deg(g2)=2
   translate(f2, n-m, g2, p );
   // g = g1 + m f2 x^(m-1) + g2 x^m
   for( i=0; i<=n-m; i++ ) g[i+m-1] = add64s(g[i+m-1],mul64s(m,f2[i],p),p);
   //for( i=0; i<=n-m-1; i++ ) g[i+m] = add64s(g[i+m],g2[i],p); 
   return;
}


int vecpartition64s(LONG *D, LONG n ) {
    LONG i,j,m;
    LONG x,t;
    m = n/2;
    x = D[m]; D[m] = D[0];
    for( i=1,j=n-1; j>=i ; ) {
       if( D[i]<x ) { D[i-1] = D[i]; i++; }
       else { t = D[i]; D[i] = D[j]; D[j] = t; j--; }
    }
    D[i-1] = x;
    return i-1;
}

void sortvec64s(LONG *D, LONG n ) {
    LONG i;
    LONG x;
    if( n<2 ) return;
    if( n>40 ) { i = vecpartition64s(D, n);
                  sortvec64s(D,i);
                  sortvec64s(D+i+1,n-i-1);
                  return;
                 }
    sortvec64s(D,n-1);
    x = D[n-1];
    for( i=n-1; i>0 && x<D[i-1]; i-- ) D[i] = D[i-1];
    D[i] = x;
    return;
}


LONG getsonly( LONG n, LONG p ) {
   int k;
   LONG s,q;
   for( s=(p-1)/2; (s&1)==0; s=s/2 );
   q = (p-1)/s;
   for( k=0; q>1; k++ ) q = q >> 1;
   while( s<2*n ) { k--; s = 2*s; }
   return s;
}


int rootsrec( LONG *lambda, LONG n, LONG *R, LONG p,
    LONG *f, LONG *g, LONG *P, LONG *Q, LONG *v, LONG *x, LONG *y, LONG *W, LONG *T,
    int TIMINGS ) {

   LONG i, k, m, okay, N;
   LONG alpha, omega, q, s;
   clock_t st,tt,ST,TT;
   recint pp;

   m = 0;
   if( n>0 && lambda[0]==0 ) { R[0] = 0; lambda++; n--; R++; m++; }
   if( n==0 ) return m;
   if( n==1 ) { R[0] = neg64s(lambda[0],p); return m+1; }

   ST = clock();

   pp = recip1(p);

   if( TIMINGS ) st = clock();
   do {
      for( i=0; i<=n; i++ ) f[i] = lambda[i];
      alpha = rand64s(p);
      changebase3( f, n, alpha, p ); // inplace
      //changebase( f, n, alpha, p ); // inplace
   } while( f[0]==0 );
   if( TIMINGS ) tt = clock();

   /* choose s parameter in the range 2n <= s < 4n */
   for( s=(p-1)/2; (s&1)==0; s=s/2 );
   q = (p-1)/s;
   for( k=0; q>1; k++ ) q = q >> 1;
   while( s<2*n ) { k--; s = 2*s; }

   printf("ROOTS :: n=%d  k=%d  s=%lld\n",n,k,s);

   if( DEBUG ) printf("alpha=%lld\n", alpha);
   if( TIMINGS ) printf("Base time = %lld ms\n",(tt-st)/1000);


if( DEBUG ) {
   polcopy64s( f, n, g );
   changebase( g, n, p-alpha, p ); // change it back
   if( n<20 ) { printf("g := "); polprint64s(g,n,p); }
   for( okay=1,i=0; i<=n; i++ ) if( g[i]!=lambda[i] ) okay = 0;
   if( okay ) printf("base okay\n"); else printf("base bug\n");
}

if( DEBUG ) {
   polcopy64s( lambda, n, g );
   if( n<20 ) { printf("lambda := "); polprint64s(g,n,p); }
   changebase( g, n, alpha, p );
   if( n<20 ) { printf("g := "); polprint64s(g,n,p); }
   for( okay=1,i=0; i<=n; i++ ) if( g[i]!=f[i] ) okay = 0;
   if( okay ) printf("base2 okay\n"); else printf("base2 bug\n");
}

   for( i=0; i<=n; i++ ) P[i] = f[i];
   //translate(f,n,Q,p);       // Translate(f,n,P,Q,p);
   poldiff64s(f,n,Q,p);       // Translate(f,n,P,Q,p);

   polcopy64s(P,n,f);
   polcopy64s(Q,n-1,g);

   //if( n<20 ) { printf("P := "); polprint64s(P,n,p); }
   //if( n<20 ) { printf("Q := "); polprint64s(Q,n-1,p); }

   for( N=8; N<=2*n; N=2*N );
   if( TIMINGS ) st = clock();
   FFTtangGraeffe64s( P, Q, n, k, T, p );
   if( TIMINGS ) tt = clock();
   if( TIMINGS ) printf("Tangent Graeffe time = %lld ms\n",(tt-st)/1000);

if( DEBUG ) {
   for( i=0; i<k; i++ ) tangGraeffe64s( f, g, n, p );
   if( !polequal64s(P,f,n) ) printf("TG: P is wrong\n"); else printf("TG: P is okay\n");
   if( !polequal64s(Q,g,n-1) ) printf("TG: Q is wrong\n"); else printf("TG: Q is okay\n");
}

   //if( n<20 ) { printf("TP := "); polprint64s(P,n,p); }
   //if( n<20 ) { printf("TQ := "); polprint64s(Q,n-1,p); }

   if( TIMINGS ) st = clock();
   omega = FFTbluestein64s( P, n, s, v, T, p );

//for( i=0; i<s; i++ ) if( v[i]==0 ) printf("Blue1: i=%d  n=%d\n",i,n);

if( DEBUG ) {
   bluestein64s( P, n, s, W, p );
   if( !vecequal64s(v,W,s) ) printf("Bluestein 1 bug\n");
   else printf("Bluestein 1 okay\n");
}

   FFTbluestein64s( Q, n-1, s, x, T, p );

if( DEBUG ) {
   bluestein64s( Q, n-1, s, W, p );
   if( !vecequal64s(x,W,s) ) printf("Bluestein 1 bug\n");
   else printf("Bluestein 2 okay\n");
}

   poldiff64s(P,n,Q,p); // derivative is in Q and has degree n-1
   FFTbluestein64s( Q, n-1, s, y, T, p );

if( DEBUG ) {
   bluestein64s( Q, n-1, s, W, p );
   if( !vecequal64s(y,W,s) ) printf("Bluestein 1 bug\n");
   else printf("Bluestein 3 okay\n");
}

   if( TIMINGS ) tt = clock();
   if( TIMINGS ) printf("Bluestein time = %lld ms\n",(tt-st)/1000);

{  LONG w,t,z,r;
   if( TIMINGS ) st = clock();
   w = 1;
   N = 0;
   m = 0;
   for( i=0; i<s; i++ ) {
       if( v[i]==0 ) N++;
       if( v[i]==0 && y[i]!=0 ) {
           t = modinv64s(y[i],p);
           //t = mul64s(t,x[i],p);
           t = mulrec64(t,x[i],pp);
           if( t==0 ) printf("DIVIDE by 0\n");
           z = modinv64s(t,p);
           r = ((LONG) 1) << k; // r = 2^k
           //r = mul64s(r,w,p);
           r = mulrec64(r,w,pp);
           //z = mul64s(r,z,p);
           z = mulrec64(r,z,pp);
           z = add64s(z,alpha,p);
           R[m++] = z;
       }
       //w = mul64s(omega,w,p);
       w = mulrec64(omega,w,pp);
   }
   if( N==0 ) { printf("No roots\n"); return 0; }
   if( TIMINGS ) tt = clock();
   if( TIMINGS ) printf("Extract roots time = %lld ms\n",(tt-st)/1000);
}

 
if( DEBUG ) if( n<50 ) { sortvec64s(R,m); printf("roots = "); vecprint64s(R,m); }

if( DEBUG ) {
   okay = 1;
   for( i=0; i<m && okay; i++ ) if( poleval64s(lambda,n,R[i],p)!=0 ) okay = 0;
   if( !okay ) printf("lambda bad root %lld\n");
}

   if( TIMINGS ) st = clock();
   fastLambda( R, m, g, p ); // Lambda( R, m, g, p );
   okay = 1;
   if( TIMINGS ) tt = clock();
   if( TIMINGS ) printf("Lambda time = %lld ms\n",(tt-st)/1000);

if( DEBUG ) {
   for( i=0; i<m && okay; i++ ) if( poleval64s(g,m,R[i],p)!=0 ) okay = 0;
   if( !okay ) { printf("roots: bad root: i=%d\n",i); return m; }
}

   if( TIMINGS ) st = clock();
   for( i=0; i<=n; i++ ) f[i] = lambda[i];
   k = FFTpoldiv64s( f, g, n, m, p );
   if( TIMINGS ) tt = clock();
   if( TIMINGS ) TT = clock();
   if( TIMINGS ) printf("divide time = %lld ms\n",(tt-st)/1000);
   if( TIMINGS ) printf("TOTAL time = %10.3f s\n",(TT-ST)/1000000.0);

   if( k!=-1 ) printf("roots: DIVIDE bad\n");

   if( n<100 ) printf("ROOTS :: %d of %d roots found\n",m,n);
   else printf("ROOTS :: %d (%4.1f%)roots found\n",m,(100.0*m)/n);

   if( m<n ) m += rootsrec( f+m, n-m, R+m, p, lambda, g, P, Q, v, x, y, W, T, 0 );
   return m;
}


int roots64s( LONG *lambda, LONG n, LONG *R, LONG p ) {
LONG m;
LONG M,N,s;
LONG *L, *f, *g, *P, *Q, *T, *v, *x, *y, *W;
   L = array64s(n+1);
   polcopy64s(lambda,n,L);
   f = array64s(n+1);
   g = array64s(n+1);
   P = array64s(n+1);
   Q = array64s(n+1);
   s = getsonly(n,p);
   printf("roots: s = %lld\n",s);
   //for( N=1; N<=2*n; N=2*N ); N = 5*N;
   for( N=1; N<=2*n; N=2*N ); N = 3*N; // eliminate B and D vectors in FFTtangGraeffe
   for( M=1; M<=2*s; M=2*M ); M = 3*M;
   printf(" N=%lld  M=%lld\n", N, M );
   if( M>N ) N = M;
   T = array64s(N);
   v = array64s(s);
   x = array64s(s);
   y = array64s(s);
   W = array64s(s);
   m = rootsrec( L, n, R, p, f, g, P, Q, v, x, y, W, T, 1 );
   free(L);
   free(T);
   free(v);
   free(x);
   free(y);
   free(P);
   free(Q);
   free(W);
   free(f);
   free(g);
   return m;
}


int compress( LONG *A, LONG n ) {
   LONG i,j,k;
   i=k=1;
   for( k=0,i=0; i<n; i=j ) {
       for( j=i+1; j<n && A[j]==A[i]; j++ ); 
       A[k++] = A[i];
   }
   return k;
}

int main( int argc, char *argv[] ) {
   LONG i,j,k,d,m,n,s,okay;
   LONG p,alpha,*A,*f,*g,*R;
   clock_t st,tt;

   seed = 1;
   mult = 6364136223846793003LL;

   n = 3*256-1;
   p = 1<<30; p = 3*p+1; // = 3 x 2^30 + 1
   p = 1<<26; p = 7*p+1; // = 7 x 2^26 + 1
   p = 1; p = p<<55; p = 5*p+1;
   p = 3*(1<<18)+1; // Allan steel
   p = 7*(1<<26)+1; // 28.8 bits for magma
   p = 6269010681299730433LL; // = 3 x 29 x 2^56 + 1
   n = 3;
   n = 3*16-1;
   n = 11;
   n = 3*1024-1;
   n = (1<<28)-1;
   n = 3*(1<<16)-1;
   n = 1000000;
   n = (1<<15)-1;
   for( i=1; i<argc; i++ ) sscanf(argv[i],"%d\n", &n);
   printf("n=%d\n",n);
   printf("p := %lld;\n",p);
   A = array64s(n);
   for( i=0; i<n; i++ ) A[i] = rand64s(p);
   for( i=0; i<n; i++ ) A[i] = i+1;
   sortvec64s(A,n); // quicksort
   k = compress(A,n); // remove duplicates
   printf( "n-k = %d duplicates\n", n-k );
   n = k;
   
   if( n<50 ) { printf("A := "); vecprint64s( A, n ); printf(";\n"); }

   f = array64s(n+1);
   g = array64s(n+1);
   st = clock();
   fastLambda( A, n, f, p );
   //Lambda( A, n, f, p );
   tt = clock()-st;
   if( n>1000 ) printf("Lambda time = %lld ms\n",tt/1000);
   if( n<20 ) { printf("f := "); polprint64s(f,n,p); }

if( DEBUG ) {
   okay = 1;
   for( i=0; i<n && okay; i++ ) if( poleval64s(f,n,A[i],p)!=0 ) { printf("BAD: i=%d\n",i); okay = 0; }
   if( okay ) printf("fast okay\n"); else printf("fast bad\n");
   //Lambda( A, n, g, p );
   //for( okay=1,i=0; i<n; i++ ) if( f[i]!=g[i] ) okay = 0; 
   //if( okay ) printf("fast okay\n"); else printf("fast bug\n");
}

   R = array64s(n);
   st = clock();
   m = roots64s( f, n, R, p );
   tt = clock();
   if( n>1000 ) printf("Total roots time = %10.3f s\n",(tt-st)/1000000.0);
   

if( DEBUG ) {
   okay = 1;
   for( i=0; i<m && okay; i++ ) if( poleval64s(f,n,R[i],p)!=0 ) { printf("BAD: i=%d\n",i); okay = 0; }
   if( okay ) printf("main: All roots okay\n"); else printf("Roots bad\n");
}
   return 1;

}