/* *************************************************************************/
/* File:   sparseldl.cc                                                    */
/* Author: Joachim Schoeberl                                               */
/* Date:   18. Nov. 2001                                                   */
/* *************************************************************************/

#include <la.hpp>
namespace ngla
{
  using namespace ngla;
  using namespace ngstd;
  

  template <class TM>
  SparseCholesky<TM> :: 
  SparseCholesky (const SparseMatrix<TM> & a, 
		  BitArray * ainner,
		  ARRAY<int> * acluster)
  { 
    //        (*testout) << "matrix = " << a << endl;
    //        (*testout) << "diag a = ";
    //        for ( int i=0; i<a.Height(); i++ ) (*testout) << i << ", " << a(i,i) << endl;

    int i, j, n;
    n = a.Height();
    height = n;

    inner = ainner;
    cluster = acluster;
    
    int printstat = 0;
    
    if (printstat)
      (cout) << "Nested Dissection: N = " << n << endl;
    
    clock_t starttime, endtime;
    starttime = clock();
    
    mdo = new MinimumDegreeOrdering (n);
    
    if (!inner && !cluster)
      for (i = 0; i < n; i++)
	for (j = 0; j < a.GetRowIndices(i).Size(); j++)
	  {
	    int col = a.GetRowIndices(i)[j];
	    if (col <= i)
	      mdo->AddEdge (i, col);
	  }

    else if (inner)
      for (i = 0; i < n; i++)
	for (j = 0; j < a.GetRowIndices(i).Size(); j++)
	  {
	    int col = a.GetRowIndices(i)[j];
	    if (col <= i)
	      if (inner->Test(i) && inner->Test(col) || i==col)
		mdo->AddEdge (i, col);
	  }

    else 
      for (i = 0; i < n; i++)
	{
	  FlatArray<const int> row = a.GetRowIndices(i);
	  for (j = 0; j < row.Size(); j++)
	    {
	      int col = row[j];
	      if (col <= i)
		if ( (*cluster)[i] == (*cluster)[col] && (*cluster)[i] ||
		     i == col )
		  mdo->AddEdge (i, col);
	    }
	}
    
    mdo->Order();

    //    (*testout) << "md-ordering, print:" << endl;
    //    mdo -> PrintCliques ();
    //    (*testout) << "order = " << mdo->order << endl;
    
    endtime = clock();
    if (printstat)
      (cout) << "ordering time = "
	     << double (endtime - starttime) / CLOCKS_PER_SEC 
	     << " secs" << endl;
    
    starttime = endtime;
    
    if (printstat)
      cout << "," << flush;
    Allocate (mdo->order, mdo->cliques, mdo->vertices);
    
    endtime = clock();
    if (printstat)
      (cout) << "allocation time = "
	     << double (endtime - starttime) / CLOCKS_PER_SEC << " secs" << endl;
    
    //  Print (cout);

    starttime = endtime;

    if (!inner && !cluster)
      for (i = 0; i < n; i++)
	for (j = 0; j < a.GetRowIndices(i).Size(); j++)
	  {
	    int col = a.GetRowIndices(i)[j];
	    if (col <= i)
	      SetOrig (i, col, a.GetRow(i)[j]);
	  }
    
    else if (inner)
      for (i = 0; i < n; i++)
	for (j = 0; j < a.GetRowIndices(i).Size(); j++)
	  {
	    int col = a.GetRowIndices(i)[j];
	    if (col <= i)
	      if (inner->Test(i) && inner->Test(col) || i==col)
		SetOrig (i, col, a.GetRow(i)[j]);
	  }
    else
      for (i = 0; i < n; i++)
	{
	  FlatArray<const int> row = a.GetRowIndices(i);
	  for (j = 0; j < row.Size(); j++)
	    {
	      int col = row[j];
	      if (col <= i)
		if ( (*cluster)[i] == (*cluster)[col] && (*cluster)[i] ||
		     i == col )
		SetOrig (i, col, a.GetRow(i)[j]);
	    }
	}
    
    if (printstat)
      cout << "do factor " << flush;
    
    Factor(&mdo->blocknr[0]);
    nze = mdo->GetNZE();

    //    delete mdo;

    if (printstat)
      cout << "done" << endl;
    
    endtime = clock();

    if (printstat)
      (cout) << " factoring time = " << double(endtime - starttime) / CLOCKS_PER_SEC << " sec" << endl;

    //           (*testout) << "cholesky factors: " << endl;
    //            Print (*testout);
  }
  
  
  template <class TM>
  SparseCholesky<TM> :: 
  SparseCholesky (const ARRAY<int> & aorder, 
		  const ARRAY<CliqueEl*> & cliques,
		  const ARRAY<MDOVertex> & vertices)
  {
    Allocate (aorder, cliques, vertices);
  }
  

  
  template <class TM>
  void SparseCholesky<TM> :: 
  Allocate (const ARRAY<int> & aorder, 
	    const ARRAY<CliqueEl*> & cliques,
	    const ARRAY<MDOVertex> & vertices)
  {
    int i, j, ii, n, cnt;
    // const CliqueEl *p1, *p2;
    
    //    SetSymmetric(1);
    
    n = aorder.Size();
    
    // order = new int[n];
    order.Alloc (n);
    order.SetName ("sparse inverse, reorder table");

    
    // int * helpi = new int[n];
    MoveableMem<int> helpi(n);
    
    // order: now inverse map 
    for (i = 0; i < n; i++)
      order[aorder[i]] = i;

    cnt = 0;
    for (i = 0; i < n; i++)
      cnt += vertices[i].nconnected;
    
    //  cout << "n = " << cnt << " =?= " << cnt2 << endl;
    
    if (n > 5000)
      cout << " " << cnt*(sizeof(TM)+sizeof(int)) << " Bytes ";
    
    // lastinrow = new int[n];
    lastinrow.Alloc(n);
    lastinrow.SetName ("sparse inverse, table 1");

    // diag = new TM[n];
    diag.Alloc (n);
    diag.SetName ("sparse inverse, diag");

    // rowindex = new int[cnt];
    rowindex.Alloc (cnt);
    rowindex.SetName ("sparse inverse, tabl 2");

    // lfact = new TM[cnt];
    lfact.Alloc (cnt);
    lfact.SetName ("sparse inverse, lfact");
    
    for (i = 0; i < cnt; i++)
      lfact[i] = 0.0;
    
    cnt = 0;
    for (i = 0; i < n; i++)
      {
	ii = aorder[i];
	for (j = 0; j < vertices[ii].nconnected; j++)
	  {
	    rowindex[cnt] = order[vertices[ii].connected[j]];
	    cnt++;
	  }
	lastinrow[i] = cnt-1;
      }
    

    
    // sort indices
    
    int first, last;
    last = -1;
    for (i = 0; i < n; i++)
      {
	first = last+1;
	last = lastinrow[i];
	
	MergeSort (last-first+1, &rowindex[first], &helpi[0]);
      }

    /*
    for (i = 0, last = -1; i < n; i++)
      {
	first = last+1;
	last = lastinrow[i];
	
	(*testout) << "row " << i << ": ";
	for (j = first; j <= last; j++)
	  (*testout) << rowindex[j] << ", ";
	(*testout) << "\n";
      }
    */

    // compressed block storage
    int nblocks = 0;
    for (i = 0, last = -1; i < n; i++)
      {
	first = last+1;
	last = lastinrow[i];

	nblocks++;
	for (j = first; j < last; j++)
	  if (rowindex[j]+1 != rowindex[j+1])
	    nblocks++;
	
      }
    
    rowindexblock.Alloc (nblocks);
    rowindexbs.Alloc (nblocks);
    lastinrowblock.Alloc (n);

    nblocks = -1;
    for (i = 0, last = -1; i < n; i++)
      {
	first = last+1;
	last = lastinrow[i];

	nblocks++;
	rowindexblock[nblocks] = first;
	rowindexbs[nblocks] = 1;

	for (j = first; j < last; j++)
	  if (rowindex[j]+1 == rowindex[j+1])
	    rowindexbs[nblocks]++;
	  else
	    {
	      nblocks++;
	      rowindexblock[nblocks] = j+1;
	      rowindexbs[nblocks] = 1;
	    }
	
	lastinrowblock[i] = nblocks;
      }


    /*
    for (i = 0, last = -1; i < n; i++)
      {
	first = last+1;
	last = lastinrowblock[i];

	(*testout) << "row " << i << ": ";
	for (j = first; j <= last; j++)
	  (*testout) << rowindexblock[j] << "-" << rowindexbs[j] << ", ";
	(*testout) << "\n";
      }
    */
    // delete helpi;
  }
  


  template <class TM>
  void SetIdentity( TM &identity )
  {
    for ( int i=0; i<identity.Width(); i++ ) identity(i,i) = 1;
  }
  void SetIdentity( double &identity ) { identity = 1; }
  void SetIdentity( Complex &identity ) { identity = 1; }


  template <class TM>
  void SparseCholesky<TM> :: 
  FactorNew (const SparseMatrix<TM> & a)
  {
    //    throw Exception ("SparseCholesky::FactorNew not implemented");
    
    if ( height != a.Height() )
      {
	cout << "SparseCholesky::FactorNew called with matrix of different size." << endl;
	return;
      }

    int i, j;
    TM id;
    id = 0.0;
    SetIdentity(id);


//     cout << "nze = " << nze << endl;
//     cout << "lastinrow[height-1] = " << lastinrow[height-1] << endl;
    for (i = 0; i < nze; i++) lfact[i] = 0.0;

    for (i = 0; i < height; i++)
      for (j = 0; j < a.GetRowIndices(i).Size(); j++)
	{
	  int col = a.GetRowIndices(i)[j];
	
	  if (!inner && !cluster || 
	      inner && inner->Test(i) && inner->Test(col) ||
	      !inner && cluster && (*cluster)[i] == (*cluster)[col] && (*cluster)[i] )
	    {
	      if ( col <= i ) SetOrig (i, col, a.GetRow(i)[j]);
	    }
	  else if (i == col)
	    SetOrig (i, i, id);
	}
    
    Factor(&mdo->blocknr[0]);
    
  }





  template <class TM>
  void SparseCholesky<TM> :: Factor (const int * blocknr)
  {
    int n, i, j, k, jj, kk, first, last;
    int firstj; //  lastj;
    TM aiinv, q, hm;
    
    //  Print (*testout);
    
    // clock_t starttime = clock();
    
    n = Height();
    if (n > 5000)
      cout << " factor " << flush;

    // to avoid aliasing:
    int * hlastinrow = lastinrow;
    int * hrowindex = rowindex;
    TM * hlfact = lfact;
    
    // int samerows = 0, diffrows = 0, sameels = 0, diffels = 0;

    last = -1;
    for (i = 0; i < n; i++)
      {
	if (n > 5000 && i % 1000 == 999)
	  {
	    if (i % 10000 == 9999)
	      cout << "+" << flush;
	    else
	      cout << "." << flush;
	  }

	first = last+1;
	last = lastinrow[i];
	
	CalcInverse (diag[i], aiinv);
	
	for (j = first; j <= last; j++)
	  {
	    jj = hrowindex[j];

	    if (blocknr[jj] == blocknr[i])
	      {
		// same indices, optimization
		if (jj > 0) firstj = hlastinrow[jj-1]+1;
		else        firstj = 0;

		TM qtrans;
		qtrans = Trans (aiinv * hlfact[j]);
		for (k = j+1; k <= last; k++, firstj++)
		  hlfact[firstj] -= qtrans * hlfact[k];

		// samerows++;
		// sameels += last-j+1;
	      }
	    else
	      {
		// have to match indices
		if (jj > 0) firstj = hlastinrow[jj-1]+1;
		else firstj = 0;
		
		TM qtrans;
		qtrans = Trans (aiinv * hlfact[j]);

		for (k = j+1; k <= last; k++)
		  {
		    kk = hrowindex[k];
		    while (hrowindex[firstj] != kk)
		      firstj++;
		    
		    hlfact[firstj] -= qtrans * hlfact[k];
		    firstj++;
		  }
	      }
	  }
	
	for (j = first; j <= last; j++)
	  {
	    q = aiinv * lfact[j];
	    diag[rowindex[j]] -= Trans (lfact[j]) * q;
	  }

	for (j = first; j <= last; j++)
	  {
	    hm = aiinv * lfact[j];
	    lfact[j] = hm;
	  }
      }
    
    /*    
    cout << "samerows: " << samerows << endl;
    cout << "diffrows: " << diffrows << endl;
    cout << "els in samerows " << sameels << endl;
    cout << "els in diffrows " << diffels << endl;
    */
    // invert diags:
    for (i = 0; i < n; i++)
      {
	CalcInverse (diag[i], hm);
	diag[i] = hm;
      }

    if (n > 5000)
      cout << endl;

    /*
      for (i = 0; i < n; i++)
      if (diag[i] == 0)
      diag[i] = 1;
    */
    //  Print (*testout);
  }
  
  


  template <class TM>
  void SparseCholesky<TM> :: 
  Mult (const BaseVector & x, BaseVector & y) const
  {
    int i, j;
    int n = Height();
    
    const FlatVector<TVX> fx = 
      dynamic_cast<const T_BaseVector<TVX> &> (x).FV();
    FlatVector<TVX> fy = 
      dynamic_cast<T_BaseVector<TVX> &> (y).FV();
    
    Vector<TVX> hy(n);
    for (i = 0; i < n; i++)
      hy(order[i]) = fx(i);

    TVX hv;

    TVX * hhy = &hy(0);

    const TM * hlfact = &lfact[0];
    const TM * hdiag = &diag[0];
    const int * hrowindex = &rowindex[0];
    const int * hlastinrow = &lastinrow[0];
    
    j = 0;
    for (i = 0; i < n; i++)
      {
	TVX val = hy(i);
	int last = hlastinrow[i];
	while (j <= last)
	  {
	    hhy[hrowindex[j]] -= Trans (hlfact[j]) * val;
	    j++;
	  }
      }
  
    for (i = 0; i < n; i++)
      {
	hv = hdiag[i] * hhy[i];
	hhy[i] = hv;
      }

    /*
    j = lastinrow[n-1];
    for (i = n-1; i >= 0; i--)
      {
	const int minj = (i > 0) ? lastinrow[i-1]+1 : 0;
	TVX sum;
	sum = 0.0;
	
	for ( ; j >= minj; --j)
	  sum += lfact[j] * hy(rowindex[j]);
	
	hy(i) -= sum;
      }
    */
    for (i = n-1; i >= 0; i--)
      {
	const int minj = (i > 0) ? lastinrow[i-1]+1 : 0;
	const int maxj = lastinrow[i];

	TVX sum;
	sum = 0.0;
	
	for (j = minj; j <= maxj; j++)
	  sum += lfact[j] * hy(rowindex[j]);
	
	hy(i) -= sum;
      }


    for (i = 0; i < n; i++)
      fy(i) = hy(order[i]);
    

    if (inner)
      {
	for (i = 0; i < n; i++)
	  if (!inner->Test(i))
	    fy(i) = 0;
      }

    if (cluster)
      {
	for (i = 0; i < n; i++)
	  if (!(*cluster)[i])
	    fy(i) = 0;
      }

  }
  
  




  template <class TM>
  void SparseCholesky<TM> :: Set (int i, int j, const TM & val)
  {
    int first, last;
  
    if (i == j)
      {
	diag[i] = val;
	return;
      }
  
    TM hval;
  
    if (i > j) 
      {
	swap (i, j);
	hval = Trans (val);
      }
    else
      hval = val;
  
    first = 0;
    if (i > 0) first = lastinrow[i-1]+1;
    last = lastinrow[i];
  
    while (first <= last)
      {
	if (rowindex[first] == j)
	  {
	    lfact[first] = hval;
	    return;
	  }
	first++;
      }
    cerr << "Position " << i << ", " << j << " not found" << endl;
  }


  template <class TM>
  const TM & SparseCholesky<TM> :: Get (int i, int j) const
  {
    int first, last;

    if (i == j)
      {
	return diag[i];
      }

    if (i > j) 
      {
	swap (i, j);
	cerr << "SparseCholesky::Get: access to upper side not available" << endl;
      }

    first = 0;
    if (i > 0) first = lastinrow[i-1]+1;
    last = lastinrow[i];
  
    while (first <= last)
      {
	if (rowindex[first] == j)
	  {
	    return lfact[first];
	  }
	first++;
      }
    cerr << "Position " << i << ", " << j << " not found" << endl;
    return *new TM;
  }

  template <class TM>
  ostream & SparseCholesky<TM> :: Print (ostream & ost) const
  {
    int i, j, n;

    n = Height();

    for (i = 0; i < n; i++)
      {
	ost << i << ": " << order[i] << " diag = "
	    << diag[i] << endl;
      }
    ost << endl;
  
    j = 1;
    for (i = 1; i <= n; i++)
      {
	ost << i << ": ";
	for ( ; j <= lastinrow[i-1]; j++)
	  {
	    ost << rowindex[j-1] << "("
		<< lfact[j] 
		<< ")  ";
	  }
	ost << endl;
      }

    return ost;
  }

  template <class TM>
  SparseCholesky<TM> :: ~SparseCholesky()
  {
    // delete [] order;
    // delete [] lastinrow;
    // delete [] rowindex;
    // delete [] lfact;
    // delete [] diag;
    delete mdo;
    delete inner;
    delete cluster;
  }

  template class SparseCholesky<double>;
  template class SparseCholesky<Complex>;
#if MAX_SYS_DIM >= 1
  template class SparseCholesky<Mat<1,1,double> >;
  template class SparseCholesky<Mat<1,1,Complex> >;
#endif
#if MAX_SYS_DIM >= 2
  template class SparseCholesky<Mat<2,2,double> >;
  template class SparseCholesky<Mat<2,2,Complex> >;
#endif
#if MAX_SYS_DIM >= 3
  template class SparseCholesky<Mat<3,3,double> >;
  template class SparseCholesky<Mat<3,3,Complex> >;
#endif
#if MAX_SYS_DIM >= 4
  template class SparseCholesky<Mat<4,4,double> >;
  template class SparseCholesky<Mat<4,4,Complex> >;
#endif
#if MAX_SYS_DIM >= 5
  template class SparseCholesky<Mat<5,5,double> >;
  template class SparseCholesky<Mat<5,5,Complex> >;
#endif
#if MAX_SYS_DIM >= 6
  template class SparseCholesky<Mat<6,6,double> >;
  template class SparseCholesky<Mat<6,6,Complex> >;
#endif
#if MAX_SYS_DIM >= 7
  template class SparseCholesky<Mat<7,7,double> >;
  template class SparseCholesky<Mat<7,7,Complex> >;
#endif
#if MAX_SYS_DIM >= 8
  template class SparseCholesky<Mat<8,8,double> >;
  template class SparseCholesky<Mat<8,8,Complex> >;
#endif
}