#include <NTL/mat_lzz_p.h>

#include <NTL/new.h>

#include <NTL/vec_long.h>
#include <NTL/vec_ulong.h>
#include <NTL/vec_double.h>

NTL_START_IMPL

NTL_matrix_impl(zz_p,vec_zz_p,vec_vec_zz_p,mat_zz_p)
NTL_io_matrix_impl(zz_p,vec_zz_p,vec_vec_zz_p,mat_zz_p)
NTL_eq_matrix_impl(zz_p,vec_zz_p,vec_vec_zz_p,mat_zz_p)


  
void add(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)   
      Error("matrix add: dimension mismatch");  
  
   X.SetDims(n, m);  
  
   long i, j;  
   for (i = 1; i <= n; i++)   
      for (j = 1; j <= m; j++)  
         add(X(i,j), A(i,j), B(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)  
      Error("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));  
}  
  


static vec_long mul_aux_vec;

static NTL_SPMM_VEC_T precon_vec;



static 
void mul_aux(mat_zz_p& X, const mat_zz_p& A, const mat_zz_p& B)  
{  
   long n = A.NumRows();  
   long l = A.NumCols();  
   long m = B.NumCols();  
  
   if (l != B.NumRows())  
      Error("matrix mul: dimension mismatch");  
  
   X.SetDims(n, m); 

   if (m > 1) {  // new preconditioning code

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

      mul_aux_vec.SetLength(m);
      long *acc = mul_aux_vec.elts();

      long i, j, k;

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

         for (j = 0; j < m; j++) acc[j] = 0;

         for (k = 0;  k < l; k++) {   
            long aa = rep(ap[k]);
            if (aa != 0) {
               const zz_p* bp = B[k].elts();
               long T1;
               mulmod_precon_t aapinv = PrepMulModPrecon(aa, p, pinv);

               for (j = 0; j < m; j++) {
        	  T1 = MulModPrecon(rep(bp[j]), aa, p, aapinv);
        	  acc[j] = AddMod(acc[j], T1, p);
               } 
            }
         }

         zz_p *xp = X[i].elts();
         for (j = 0; j < m; j++)
            xp[j].LoopHole() = acc[j];    
      }

   }
   else {  // just use the old code, w/o preconditioning

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

      long i, j, k;  
      long acc, tmp;  

      for (i = 1; i <= n; i++) {  
	 for (j = 1; j <= m; j++) {  
            acc = 0;  
            for(k = 1; k <= l; k++) {  
               tmp = MulMod(rep(A(i,k)), rep(B(k,j)), p, pinv);  
               acc = AddMod(acc, tmp, p);  
            }  
            X(i,j).LoopHole() = acc;  
	 } 
      }
  
   }
}  

void mul(mat_zz_p& X, const mat_zz_p& A, const mat_zz_p& B)  
{  
   if (&X == &A || &X == &B) {  
      mat_zz_p tmp;  
      mul_aux(tmp, A, B);  
      X = tmp;  
   }  
   else  
      mul_aux(X, A, B);  
}  


void mul(vec_zz_p& x, const vec_zz_p& a, const mat_zz_p& B)
{
   long l = a.length();
   long m = B.NumCols();
  
   if (l != B.NumRows())  
      Error("matrix mul: dimension mismatch");  

   if (m == 0) { 

      x.SetLength(0);
      
   }
   else if (m == 1) {

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

      long acc, tmp;
      long k;

      acc = 0;  
      for(k = 1; k <= l; k++) {  
         tmp = MulMod(rep(a(k)), rep(B(k,1)), p, pinv);  
         acc = AddMod(acc, tmp, p);  
      } 

      x.SetLength(1);
      x(1).LoopHole()  = acc;
          
   }
   else {  // m > 1.  precondition


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

      mul_aux_vec.SetLength(m);
      long *acc = mul_aux_vec.elts();

      long j, k;


      const zz_p* ap = a.elts();

      for (j = 0; j < m; j++) acc[j] = 0;

      for (k = 0;  k < l; k++) {
         long aa = rep(ap[k]);
         if (aa != 0) {
            const zz_p* bp = B[k].elts();
            long T1;
            mulmod_precon_t aapinv = PrepMulModPrecon(aa, p, pinv);

            for (j = 0; j < m; j++) {
               T1 = MulModPrecon(rep(bp[j]), aa, p, aapinv);
               acc[j] = AddMod(acc[j], T1, p);
            }
         } 
      }

      x.SetLength(m);
      zz_p *xp = x.elts();
      for (j = 0; j < m; j++)
         xp[j].LoopHole() = acc[j];    
      
   }
}

  
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())
      Error("matrix mul: dimension mismatch");

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

   long p = zz_p::modulus();
   double 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 {

      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 mul(vec_zz_p& x, const mat_zz_p& A, const vec_zz_p& b)  
{  
   if (&b == &x || A.position1(x) != -1) {
      vec_zz_p tmp;
      mul_aux(tmp, A, b);
      x = tmp;
   }
   else
      mul_aux(x, A, b);

}  


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();
      double 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 mul(mat_zz_p& X, const mat_zz_p& A, long b_in)
{
   zz_p b;
   b = b_in;
   mul(X, A, b);
} 




     
  
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));  
} 



void determinant(zz_p& d, const mat_zz_p& M_in)
{
   long k, n;
   long i, j;
   long pos;
   zz_p t1, t2, t3;
   zz_p *x, *y;

   mat_zz_p M;
   M = M_in;

   n = M.NumRows();

   if (M.NumCols() != n)
      Error("determinant: nonsquare matrix");

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

   zz_p det;

   set(det);

   long p = zz_p::modulus();
   double 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]);

         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 < 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;
      }
   }

   d = det;
}




long IsIdent(const mat_zz_p& A, long n)
{
   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 (!IsOne(A(i, j))) return 0;
         }

   return 1;
}
            

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 solve(zz_p& d, vec_zz_p& X, 
           const mat_zz_p& A, const vec_zz_p& b)

{
   long n = A.NumRows();

   if (A.NumCols() != n)
      Error("solve: nonsquare matrix");