void power(mat_zz_p& X, const mat_zz_p& A, const ZZ& e)
{
   if (A.NumRows() != A.NumCols()) LogicError("power: non-square matrix");

   if (e == 0) {
      ident(X, A.NumRows());
      return;
   }

   mat_zz_p T1, T2;
   long i, k;

   k = NumBits(e);
   T1 = A;

   for (i = k-2; i >= 0; i--) {
      sqr(T2, T1);
      if (bit(e, i))
         mul(T1, T2, A);
      else
         T1 = T2;
   }

   if (e < 0)
      inv(X, T1);
   else
      X = T1;
}

void transpose(mat_zz_p& X, const mat_zz_p& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   long i, j;

   if (&X == & A) {
      if (n == m)
         for (i = 1; i <= n; i++)
            for (j = i+1; j <= n; j++)
               swap(X(i, j), X(j, i));
      else {
         mat_zz_p tmp;
         tmp.SetDims(m, n);
         for (i = 1; i <= n; i++)
            for (j = 1; j <= m; j++)
               tmp(j, i) = A(i, j);
         X.kill();
         X = tmp;
      }
   }
   else {
      X.SetDims(m, n);
      for (i = 1; i <= n; i++)
         for (j = 1; j <= m; j++)
            X(j, i) = A(i, j);
   }
}

void mul(mat_zz_p& X, const mat_zz_p& A, zz_p b)
{
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;

   if (n == 0 || m == 0 || (n == 1 && m == 1)) {

      for (i = 0; i < n; i++)
	 for (j = 0; j < m; j++)
            mul(X[i][j], A[i][j], b);

   }
   else {
      
      long p = zz_p::modulus();
      mulmod_t pinv = zz_p::ModulusInverse();
      long bb = rep(b);
      mulmod_precon_t bpinv = PrepMulModPrecon(bb, p, pinv);
      
      for (i = 0; i < n; i++) {
         const zz_p *ap = A[i].elts();
         zz_p *xp = X[i].elts();

	 for (j = 0; j < m; j++)
            xp[j].LoopHole() = MulModPrecon(rep(ap[j]), bb, p, bpinv);
      }

   }
}

void kernel(mat_zz_p& X, const mat_zz_p& A)
{
   long m = A.NumRows();
   long n = A.NumCols();

   mat_zz_p M;
   long r;

   transpose(M, A);
   r = gauss(M);

   X.SetDims(m-r, m);

   long i, j, k, s;
   zz_p t1, t2;

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

   vec_zz_p inverses;
   inverses.SetLength(m);

   j = -1;
   for (i = 0; i < r; i++) {
      do {
         j++;
      } while (IsZero(M[i][j]));

      D[j] = i;
      inv(inverses[j], M[i][j]); 
   }

   for (k = 0; k < m-r; k++) {
      vec_zz_p& v = X[k];
      long pos = 0;
      for (j = m-1; j >= 0; j--) {
         if (D[j] == -1) {
            if (pos == k)
               set(v[j]);
            else
               clear(v[j]);
            pos++;
         }
         else {
            i = D[j];

            clear(t1);

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

            mul(t1, t1, inverses[j]);
            negate(v[j], t1);
         }
      }
   }
}

void mul_aux(vec_zz_p& x, const mat_zz_p& A, const vec_zz_p& b)
{
   long n = A.NumRows();
   long l = A.NumCols();

   if (l != b.length())
      LogicError("matrix mul: dimension mismatch");

   x.SetLength(n);
   zz_p* xp = x.elts();

   long p = zz_p::modulus();
   mulmod_t pinv = zz_p::ModulusInverse();

   long i, k;
   long acc, tmp;

   const zz_p* bp = b.elts();

   if (n <= 1) {

      for (i = 0; i < n; i++) {
	 acc = 0;
	 const zz_p* ap = A[i].elts();

	 for (k = 0; k < l; k++) {
            tmp = MulMod(rep(ap[k]), rep(bp[k]), p, pinv);
            acc = AddMod(acc, tmp, p);
	 }

	 xp[i].LoopHole() = acc;
      }

   }
   else {

      Vec<mulmod_precon_t>::Watcher watch_precon_vec(precon_vec);
      precon_vec.SetLength(l);
      mulmod_precon_t *bpinv = precon_vec.elts();

      for (k = 0; k < l; k++)
         bpinv[k] = PrepMulModPrecon(rep(bp[k]), p, pinv);

      for (i = 0; i < n; i++) {
	 acc = 0;
	 const zz_p* ap = A[i].elts();

	 for (k = 0; k < l; k++) {
            tmp = MulModPrecon(rep(ap[k]), rep(bp[k]), p, bpinv[k]);
            acc = AddMod(acc, tmp, p);
	 }

	 xp[i].LoopHole() = acc;
      } 
   }
}

void negate(mat_zz_p& X, const mat_zz_p& A)
{
   long n = A.NumRows();
   long m = A.NumCols();


   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         negate(X(i,j), A(i,j));
}

void cvt(mat_GF2& x, const mat_zz_p& a)
{
    long n = a.NumRows();
    long m = a.NumCols();

    x.SetDims(n, m);

    long i, j;

    for (i = 0; i < n; i++)
        for (j = 0; j < m; j++)
            x.put(i, j, rep(a[i][j]));
}

void sub(mat_zz_p& X, const mat_zz_p& A, const mat_zz_p& B)  
{  
   long n = A.NumRows();  
   long m = A.NumCols();  
  
   if (B.NumRows() != n || B.NumCols() != m)  
      LogicError("matrix sub: dimension mismatch");  
  
   X.SetDims(n, m);  
  
   long i, j;  
   for (i = 1; i <= n; i++)  
      for (j = 1; j <= m; j++)  
         sub(X(i,j), A(i,j), B(i,j));  
}  

void clear(mat_zz_p& x)
{
   long n = x.NumRows();
   long i;
   for (i = 0; i < n; i++)
      clear(x[i]);
}

long IsDiag(const mat_zz_p& A, long n, zz_p d)
{
   if (A.NumRows() != n || A.NumCols() != n)
      return 0;

   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i != j) {
            if (!IsZero(A(i, j))) return 0;
         }
         else {
            if (A(i, j) != d) return 0;
         }

   return 1;
}

void conv(mat_zz_p& x, const mat_ZZ& a)
{
    long n = a.NumRows();
    long m = a.NumCols();
    long i;

    x.SetDims(n, m);
    for (i = 0; i < n; i++)
        conv(x[i], a[i]);
}

long IsZero(const mat_zz_p& a)
{
   long n = a.NumRows();
   long i;

   for (i = 0; i < n; i++)
      if (!IsZero(a[i]))
         return 0;

   return 1;
}

void diag(mat_zz_p& X, long n, zz_p d)  
{  
   X.SetDims(n, n);  
   long i, j;  
  
   for (i = 1; i <= n; i++)  
      for (j = 1; j <= n; j++)  
         if (i == j)  
            X(i, j) = d;  
         else  
            clear(X(i, j));  
} 

NTL_CLIENT


void random(mat_zz_p& X, long n, long m)
{
    X.SetDims(n, m);
    long i, j;

    for (i = 0; i < n; i++)
        for (j = 0; j < m; j++)
            random(X[i][j]);
}

void ident(mat_zz_p& X, long n)  
{  
   X.SetDims(n, n);  
   long i, j;  
  
   for (i = 1; i <= n; i++)  
      for (j = 1; j <= n; j++)  
         if (i == j)  
            set(X(i, j));  
         else  
            clear(X(i, j));  
} 

// FIXME: at some point need to make a template for these two functions
// prime power solver
// A is an n x n matrix, we compute its inverse mod p^r. An error is raised
// if A is not inverible mod p. zz_p::modulus() is assumed to be p^r, for
// p prime, r >= 1.
void ppInvert(mat_zz_p& X, const mat_zz_p& A, long p, long r)
{
  if (r == 1) { // use native inversion from NTL
    inv(X, A);    // X = A^{-1}
    return;
  }

  // begin by inverting A modulo p

  // convert to long for a safe transaltion to mod-p objects
  Mat<long> tmp;
  conv(tmp, A);
  { // open a new block for mod-p computation
  zz_pBak bak_pr; bak_pr.save(); // backup the mod-p^r moduli
  zz_p::init(p);   // Set the mod-p moduli

  mat_zz_p A1, Inv1;
  conv(A1, tmp);   // Recover A as a mat_zz_pE object modulo p
  inv(Inv1, A1);      // Inv1 = A^{-1} (mod p)
  conv(tmp, Inv1); // convert to long for transaltion to a mod-p^r object
  } // mod-p^r moduli restored on desctuction of bak_pr and bak_prE
  mat_zz_p XX;
  conv(XX, tmp); // XX = A^{-1} (mod p)

  // Now lift the solution modulo p^r

  // Compute the "correction factor" Z, s.t. XX*A = I - p*Z (mod p^r)
  long n = A.NumRows();
  const mat_zz_p I = ident_mat_zz_p(n); // identity matrix
  mat_zz_p Z = I - XX*A;

  conv(tmp, Z);  // Conver to long to divide by p
  for (long i=0; i<n; i++) for (long j=0; j<n; j++) tmp[i][j] /= p;
  conv(Z, tmp);  // convert back to a mod-p^r object

  // The inverse of A is ( I+(pZ)+(pZ)^2+...+(pZ)^{r-1} )*XX (mod p^r). We use
  // O(log r) products to copmute it as (I+pZ)* (I+(pZ)^2)* (I+(pZ)^4)*...* XX

  long e = NextPowerOfTwo(r); // 2^e is smallest power of two >= r

  Z *= p;                 // = pZ
  mat_zz_p prod = I + Z; // = I + pZ
  for (long i=1; i<e; i++) {
    sqr(Z, Z);     // = (pZ)^{2^i}
    prod *= (I+Z); // = sum_{j=0}^{2^{i+1}-1} (pZ)^j
  }
  mul(X, prod, XX); // X = A^{-1} mod p^r
  assert(X*A == I);
}

long CRT(mat_ZZ& gg, ZZ& a, const mat_zz_p& G)
{
    long n = gg.NumRows();
    long m = gg.NumCols();

    if (G.NumRows() != n || G.NumCols() != m)
        Error("CRT: dimension mismatch");

    long p = zz_p::modulus();

    ZZ new_a;
    mul(new_a, a, p);

    long a_inv;
    a_inv = rem(a, p);
    a_inv = InvMod(a_inv, p);

    long p1;
    p1 = p >> 1;

    ZZ a1;
    RightShift(a1, a, 1);

    long p_odd = (p & 1);

    long modified = 0;

    long h;

    ZZ g;
    long i, j;

    for (i = 0; i < n; i++) {
        for (j = 0; j < m; j++) {
            if (!CRTInRange(gg[i][j], a)) {
                modified = 1;
                rem(g, gg[i][j], a);
                if (g > a1) sub(g, g, a);
            }
            else
                g = gg[i][j];

            h = rem(g, p);
            h = SubMod(rep(G[i][j]), h, p);
            h = MulMod(h, a_inv, p);
            if (h > p1)
                h = h - p;

            if (h != 0) {
                modified = 1;

                if (!p_odd && g > 0 && (h == p1))
                    MulSubFrom(g, a, h);
                else
                    MulAddTo(g, a, h);

            }

            gg[i][j] = g;
        }
    }

    a = new_a;

    return modified;

}

NTL_START_IMPL


void CharPoly(zz_pX& f, const mat_zz_p& M)
{
   long n = M.NumRows();
   if (M.NumCols() != n)
      LogicError("CharPoly: nonsquare matrix");

   if (n == 0) {
      set(f);
      return;
   }

   zz_p t;

   if (n == 1) {
      SetX(f);
      negate(t, M(1, 1));
      SetCoeff(f, 0, t);
      return;
   }

   mat_zz_p H;

   H = M;

   long i, j, m;
   zz_p u, t1;

   for (m = 2; m <= n-1; m++) {
      i = m;
      while (i <= n && IsZero(H(i, m-1)))
         i++;

      if (i <= n) {
         t = H(i, m-1);
         if (i > m) {
            swap(H(i), H(m));
            // swap columns i and m
            for (j = 1; j <= n; j++) 
               swap(H(j, i), H(j, m));
         }

         for (i = m+1; i <= n; i++) {
            div(u, H(i, m-1), t);
            for (j = m; j <= n; j++) {
               mul(t1, u, H(m, j));
               sub(H(i, j), H(i, j), t1);
            }

            for (j = 1; j <= n; j++) {
               mul(t1, u, H(j, i));
               add(H(j, m), H(j, m), t1);
            }
         }
      }
   }

   vec_zz_pX F;
   F.SetLength(n+1);
   zz_pX T;
   T.SetMaxLength(n);

   set(F[0]);
   for (m = 1; m <= n; m++) {
      LeftShift(F[m], F[m-1], 1);
      mul(T, F[m-1], H(m, m));
      sub(F[m], F[m], T);

      set(t);
      for (i = 1; i <= m-1; i++) {
         mul(t, t, H(m-i+1, m-i));
         mul(t1, t, H(m-i, m));
         mul(T, F[m-i-1], t1);
         sub(F[m], F[m], T);
      }
   }

   f = F[n];
}

void inv(zz_p& d, mat_zz_p& X, const mat_zz_p& A)
{
   long n = A.NumRows();
   if (A.NumCols() != n)
      LogicError("inv: nonsquare matrix");

   if (n == 0) {
      set(d);
      X.SetDims(0, 0);
      return;
   }

   long i, j, k, pos;
   zz_p t1, t2, t3;
   zz_p *x, *y;

   mat_zz_p M;
   M.SetDims(n, 2*n);
   for (i = 0; i < n; i++) {
      for (j = 0; j < n; j++) {
         M[i][j] = A[i][j];
         clear(M[i][n+j]);
      }
      set(M[i][n+i]);
   }

   zz_p det;
   set(det);

   long p = zz_p::modulus();
   mulmod_t pinv = zz_p::ModulusInverse();

   for (k = 0; k < n; k++) {
      pos = -1;
      for (i = k; i < n; i++) {
         if (!IsZero(M[i][k])) {
            pos = i;
            break;
         }
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            negate(det, det);
         }

         mul(det, det, M[k][k]);

         inv(t3, M[k][k]);
         M[k][k] = t3;

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

            mul(t1, M[i][k], t3);
            negate(t1, t1);

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

            long T1 = rep(t1);
            mulmod_precon_t t1pinv = PrepMulModPrecon(T1, p, pinv); // T1*pinv;
            long T2;

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

               T2 = MulModPrecon(rep(*y), T1, p, t1pinv);
               x->LoopHole() = AddMod(rep(*x), T2, p);
            }
         }
      }
      else {
         clear(d);
         return;
      }
   }

   X.SetDims(n, n);
   for (k = 0; k < n; k++) {
      for (i = n-1; i >= 0; i--) {
         clear(t1);
         for (j = i+1; j < n; j++) {
            mul(t2, X[j][k], M[i][j]);
            add(t1, t1, t2);
         }
         sub(t1, M[i][n+k], t1);
         mul(X[i][k], t1, M[i][i]);
      }
   }

   d = det;
}

long gauss(mat_zz_p& M, long w)
{
   long k, l;
   long i, j;
   long pos;
   zz_p t1, t2, t3;
   zz_p *x, *y;

   long n = M.NumRows();
   long m = M.NumCols();

   if (w < 0 || w > m)
      LogicError("gauss: bad args");

   long p = zz_p::modulus();
   mulmod_t pinv = zz_p::ModulusInverse();
   long T1, T2;

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

      pos = -1;
      for (i = l; i < n; i++) {
         if (!IsZero(M[i][k])) {
            pos = i;
            break;
         }
      }

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

         inv(t3, M[l][k]);
         negate(t3, t3);

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

            mul(t1, M[i][k], t3);

            T1 = rep(t1);
            mulmod_precon_t T1pinv = PrepMulModPrecon(T1, p, pinv); 

            clear(M[i][k]);

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

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

               T2 = MulModPrecon(rep(*y), T1, p, T1pinv);
               T2 = AddMod(T2, rep(*x), p);
               (*x).LoopHole() = T2;
            }
         }

         l++;
      }
   }

   return l;
}