#include <NTL/GF2EXFactoring.h>
#include <NTL/vec_GF2XVec.h>
#include <NTL/fileio.h>
#include <NTL/FacVec.h>

#include <stdio.h>

#include <NTL/new.h>

NTL_START_IMPL




static
void IterSqr(GF2E& c, const GF2E& a, long n)
{
   GF2E res;

   long i;

   res = a;

   for (i = 0; i < n; i++)
      sqr(res, res);

   c = res;
}
   


void SquareFreeDecomp(vec_pair_GF2EX_long& u, const GF2EX& ff)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("SquareFreeDecomp: bad args");

   GF2EX r, t, v, tmp1;
   long m, j, finished, done;

   u.SetLength(0);

   if (deg(f) == 0)
      return;

   m = 1;
   finished = 0;

   do {
      j = 1;
      diff(tmp1, f);
      GCD(r, f, tmp1);
      div(t, f, r);

      if (deg(t) > 0) {
         done = 0;
         do {
            GCD(v, r, t);
            div(tmp1, t, v);
            if (deg(tmp1) > 0) append(u, cons(tmp1, j*m));
            if (deg(v) > 0) {
               div(r, r, v);
               t = v;
               j++;
            }
            else
               done = 1;
         } while (!done);
         if (deg(r) == 0) finished = 1;
      }

      if (!finished) {
         /* r is a square */

         long k, d;
         d = deg(r)/2;
         f.rep.SetLength(d+1);
         for (k = 0; k <= d; k++) 
            IterSqr(f.rep[k], r.rep[k*2], GF2E::degree()-1);
         m = m*2;
      }
   } while (!finished);
}
         


static
void NullSpace(long& r, vec_long& D, vec_GF2XVec& M, long verbose)
{
   long k, l, n;
   long i, j;
   long pos;
   GF2X t1, t2;
   GF2X *x, *y;

   const GF2XModulus& p = GF2E::modulus();

   n = M.length();

   D.SetLength(n);
   for (j = 0; j < n; j++) D[j] = -1;

   r = 0;

   l = 0;
   for (k = 0; k < n; k++) {

      if (verbose && k % 10 == 0) cerr << "+";

      pos = -1;
      for (i = l; i < n; i++) {
         rem(t1, M[i][k], p);
         M[i][k] = t1;
         if (pos == -1 && !IsZero(t1))
            pos = i;
      }

      if (pos != -1) {
         swap(M[pos], M[l]);

         // make M[l, k] == -1 mod p, and make row l reduced

         InvMod(t1, M[l][k], p);
         for (j = k+1; j < n; j++) {
            rem(t2, M[l][j], p);
            MulMod(M[l][j], t2, t1, p);
         }

         for (i = l+1; i < n; i++) {
            // M[i] = M[i] + M[l]*M[i,k]

            t1 = M[i][k];   // this is already reduced

            x = M[i].elts() + (k+1);
            y = M[l].elts() + (k+1);

            for (j = k+1; j < n; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(*x, *x, t2);
            }
         }

         D[k] = l;   // variable k is defined by row l
         l++;

      }
      else {
         r++;
      }
   }
}



static
void BuildMatrix(vec_GF2XVec& M, long n, const GF2EX& g, const GF2EXModulus& F,
                 long verbose)
{
   long i, j, m;
   GF2EX h;


   M.SetLength(n);
   for (i = 0; i < n; i++)
      M[i].SetSize(n, 2*GF2E::WordLength());

   set(h);
   for (j = 0; j < n; j++) {
      if (verbose && j % 10 == 0) cerr << "+";

      m = deg(h);
      for (i = 0; i < n; i++) {
         if (i <= m)
            M[i][j] = rep(h.rep[i]);
         else
            clear(M[i][j]);
      }

      if (j < n-1)
         MulMod(h, h, g, F);
   }

   for (i = 0; i < n; i++)
      add(M[i][i], M[i][i], 1);

}


static
void TraceMap(GF2EX& h, const GF2EX& a, const GF2EXModulus& F)

// one could consider making a version based on modular composition,
// as in ComposeFrobeniusMap...

{
   GF2EX res, tmp;

   res = a;
   tmp = a;

   long i;
   for (i = 0; i < GF2E::degree()-1; i++) {
      SqrMod(tmp, tmp, F);
      add(res, res, tmp);
   }

   h = res;
}

void PlainFrobeniusMap(GF2EX& h, const GF2EXModulus& F)
{
   GF2EX res;

   SetX(res);
   long i;
   for (i = 0; i < GF2E::degree(); i++) 
      SqrMod(res, res, F);

   h = res;
}

long UseComposeFrobenius(long d, long n)
{
   long i;
   i = 1;
   while (i <= d) i = i << 1;
   i = i >> 1;

   i = i >> 1;
   long m = 1;

   long dz;

   if (n == 2) {
      dz = 1;
   }
   else {
      while (i) {
         long m1 = 2*m;
         if (i & d) m1++;
   
         if (m1 >= NTL_BITS_PER_LONG-1 || (1L << m1) >= n) break;
   
         m = m1;
         i = i >> 1;
      }

      dz = 1L << m;
   }

   long rootn = SqrRoot(n);
   long cnt = 0;

   if (i) {
      cnt += SqrRoot(dz+1);
      i = i >> 1;
   }

   while (i) {
      cnt += rootn;
      i = i >> 1;
   }

   return 4*cnt <= d;
}

void ComposeFrobeniusMap(GF2EX& y, const GF2EXModulus& F)
{
   long d = GF2E::degree();
   long n = deg(F);

   long i;
   i = 1;
   while (i <= d) i = i << 1;
   i = i >> 1;

   GF2EX z(INIT_SIZE, n), z1(INIT_SIZE, n);

   i = i >> 1;
   long m = 1;

   if (n == 2) {
      SetX(z);
      SqrMod(z, z, F);
   }
   else {
      while (i) {
         long m1 = 2*m;
         if (i & d) m1++;
   
         if (m1 >= NTL_BITS_PER_LONG-1 || (1L << m1) >= n) break;
   
         m = m1;
         i = i >> 1;
      }

      clear(z);
      SetCoeff(z, 1L << m);
   }


   while (i) {
      z1 = z;

      long j, k, dz;
      dz = deg(z);

      for (j = 0; j <= dz; j++)
         for (k = 0; k < m; k++)
            sqr(z1.rep[j], z1.rep[j]);

      CompMod(z, z1, z, F);
      m = 2*m;

      if (d & i) {
         SqrMod(z, z, F);
         m++;
      }

      i = i >> 1;
   }

   y = z;
}

void FrobeniusMap(GF2EX& h, const GF2EXModulus& F)
{
   long n = deg(F);
   long d = GF2E::degree();

   if (n == 1) {
      h = ConstTerm(F);
      return;
   }

   if (UseComposeFrobenius(d, n))
      ComposeFrobeniusMap(h, F);
   else
      PlainFrobeniusMap(h, F);
}



   



static
void RecFindRoots(vec_GF2E& x, const GF2EX& f)
{
   if (deg(f) == 0) return;

   if (deg(f) == 1) {
      long k = x.length();
      x.SetLength(k+1);
      x[k] = ConstTerm(f);
      return;
   }
      
   GF2EX h;

   GF2E r;

   
   {
      GF2EXModulus F;
      build(F, f);

      do {
         random(r);
         clear(h);
         SetCoeff(h, 1, r);
         TraceMap(h, h, F);
         GCD(h, h, f);
      } while (deg(h) <= 0 || deg(h) == deg(f));
   }

   RecFindRoots(x, h);
   div(h, f, h); 
   RecFindRoots(x, h);
}

void FindRoots(vec_GF2E& x, const GF2EX& ff)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("FindRoots: bad args");

   x.SetMaxLength(deg(f));
   x.SetLength(0);
   RecFindRoots(x, f);
}


static
void RandomBasisElt(GF2EX& g, const vec_long& D, const vec_GF2XVec& M)
{
   static GF2X t1, t2;

   long n = D.length();

   long i, j, s;

   g.rep.SetLength(n);

   vec_GF2E& v = g.rep;

   for (j = n-1; j >= 0; j--) {
      if (D[j] == -1)
         random(v[j]);
      else {
         i = D[j];

         // v[j] = sum_{s=j+1}^{n-1} v[s]*M[i,s]

         clear(t1);

         for (s = j+1; s < n; s++) {
            mul(t2, rep(v[s]), M[i][s]);
            add(t1, t1, t2);
         }

         conv(v[j], t1);
      }
   }

   g.normalize();
}



static
void split(GF2EX& f1, GF2EX& g1, GF2EX& f2, GF2EX& g2,
           const GF2EX& f, const GF2EX& g, 
           const vec_GF2E& roots, long lo, long mid)
{
   long r = mid-lo+1;

   GF2EXModulus F;
   build(F, f);

   vec_GF2E lroots(INIT_SIZE, r);
   long i;

   for (i = 0; i < r; i++)
      lroots[i] = roots[lo+i];


   GF2EX h, a, d;
   BuildFromRoots(h, lroots);
   CompMod(a, h, g, F);


   GCD(f1, a, f);
   
   div(f2, f, f1);

   rem(g1, g, f1);
   rem(g2, g, f2);
}

static
void RecFindFactors(vec_GF2EX& factors, const GF2EX& f, const GF2EX& g,
                    const vec_GF2E& roots, long lo, long hi)
{
   long r = hi-lo+1;

   if (r == 0) return;

   if (r == 1) {
      append(factors, f);
      return;
   }

   GF2EX f1, g1, f2, g2;

   long mid = (lo+hi)/2;

   split(f1, g1, f2, g2, f, g, roots, lo, mid);

   RecFindFactors(factors, f1, g1, roots, lo, mid);
   RecFindFactors(factors, f2, g2, roots, mid+1, hi);
}


static
void FindFactors(vec_GF2EX& factors, const GF2EX& f, const GF2EX& g,
                 const vec_GF2E& roots)
{
   long r = roots.length();

   factors.SetMaxLength(r);
   factors.SetLength(0);

   RecFindFactors(factors, f, g, roots, 0, r-1);
}

#if 0

static
void IterFindFactors(vec_GF2EX& factors, const GF2EX& f,
                     const GF2EX& g, const vec_GF2E& roots)
{
   long r = roots.length();
   long i;
   GF2EX h;

   factors.SetLength(r);

   for (i = 0; i < r; i++) {
      add(h, g, roots[i]);
      GCD(factors[i], f, h);
   }
}

#endif


   

void SFBerlekamp(vec_GF2EX& factors, const GF2EX& ff, long verbose)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("SFBerlekamp: bad args");

   if (deg(f) == 0) {
      factors.SetLength(0);
      return;
   }

   if (deg(f) == 1) {
      factors.SetLength(1);
      factors[0] = f;
      return;
   }

   double t;

   long n = deg(f);

   GF2EXModulus F;

   build(F, f);

   GF2EX g, h;

   if (verbose) { cerr << "computing X^p..."; t = GetTime(); }
   FrobeniusMap(g, F);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   vec_long D;
   long r;

   vec_GF2XVec M;

   if (verbose) { cerr << "building matrix..."; t = GetTime(); }
   BuildMatrix(M, n, g, F, verbose);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   if (verbose) { cerr << "diagonalizing..."; t = GetTime(); }
   NullSpace(r, D, M, verbose);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }


   if (verbose) cerr << "number of factors = " << r << "\n";

   if (r == 1) {
      factors.SetLength(1);
      factors[0] = f;
      return;
   }

   if (verbose) { cerr << "factor extraction..."; t = GetTime(); }

   vec_GF2E roots;

   RandomBasisElt(g, D, M);
   MinPolyMod(h, g, F, r);
   if (deg(h) == r) M.kill();
   FindRoots(roots, h);
   FindFactors(factors, f, g, roots);

   GF2EX g1;
   vec_GF2EX S, S1;
   long i;

   while (factors.length() < r) {
      if (verbose) cerr << "+";
      RandomBasisElt(g, D, M);
      S.kill();
      for (i = 0; i < factors.length(); i++) {
         const GF2EX& f = factors[i];
         if (deg(f) == 1) {
            append(S, f);
            continue;
         }
         build(F, f);
         rem(g1, g, F);
         if (deg(g1) <= 0) {
            append(S, f);
            continue;
         }
         MinPolyMod(h, g1, F, min(deg(f), r-factors.length()+1));
         FindRoots(roots, h);
         S1.kill();
         FindFactors(S1, f, g1, roots);
         append(S, S1);
      }
      swap(factors, S);
   }

   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   if (verbose) {
      cerr << "degrees:";
      long i;
      for (i = 0; i < factors.length(); i++)
         cerr << " " << deg(factors[i]);
      cerr << "\n";
   }
}


void berlekamp(vec_pair_GF2EX_long& factors, const GF2EX& f, long verbose)
{
   double t;
   vec_pair_GF2EX_long sfd;
   vec_GF2EX x;

   if (!IsOne(LeadCoeff(f)))
      Error("berlekamp: bad args");

   
   if (verbose) { cerr << "square-free decomposition..."; t = GetTime(); }
   SquareFreeDecomp(sfd, f);
   if (verbose) cerr << (GetTime()-t) << "\n";

   factors.SetLength(0);

   long i, j;

   for (i = 0; i < sfd.length(); i++) {
      if (verbose) {
         cerr << "factoring multiplicity " << sfd[i].b 
              << ", deg = " << deg(sfd[i].a) << "\n";
      }

      SFBerlekamp(x, sfd[i].a, verbose);

      for (j = 0; j < x.length(); j++)
         append(factors, cons(x[j], sfd[i].b));
   }
}



static
void AddFactor(vec_pair_GF2EX_long& factors, const GF2EX& g, long d, long verbose)
{
   if (verbose)
      cerr << "degree=" << d << ", number=" << deg(g)/d << "\n";
   append(factors, cons(g, d));
}

static
void ProcessTable(GF2EX& f, vec_pair_GF2EX_long& factors, 
                  const GF2EXModulus& F, long limit, const vec_GF2EX& tbl,
                  long d, long verbose)

{
   if (limit == 0) return;

   if (verbose) cerr << "+";

   GF2EX t1;

   if (limit == 1) {
      GCD(t1, f, tbl[0]);
      if (deg(t1) > 0) {
         AddFactor(factors, t1, d, verbose);
         div(f, f, t1);
      }

      return;
   }

   long i;

   t1 = tbl[0];
   for (i = 1; i < limit; i++)
      MulMod(t1, t1, tbl[i], F);

   GCD(t1, f, t1);

   if (deg(t1) == 0) return;

   div(f, f, t1);

   GF2EX t2;

   i = 0;
   d = d - limit + 1;

   while (2*d <= deg(t1)) {
      GCD(t2, tbl[i], t1); 
      if (deg(t2) > 0) {
         AddFactor(factors, t2, d, verbose);
         div(t1, t1, t2);
      }

      i++;
      d++;
   }

   if (deg(t1) > 0)
      AddFactor(factors, t1, deg(t1), verbose);
}


void TraceMap(GF2EX& w, const GF2EX& a, long d, const GF2EXModulus& F, 
              const GF2EX& b)

{