/*------------------------------------------------------------------------------*/
/** 
 *  \file   GW_SparseMatrix.h
 *  \brief  Definition of class \c GW_SparseMatrix
 *  \author Gabriel Peyr
 *  \date   6-8-2003
 */ 
/*------------------------------------------------------------------------------*/

#ifndef _GW_SPARSEMATRIX_H_
#define _GW_SPARSEMATRIX_H_

#include "GW_MathsConfig.h"


extern "C" {
#include <laspack/errhandl.h>
#include <laspack/vector.h>
#include <laspack/qmatrix.h>
#include <laspack/operats.h>
#include <laspack/precond.h>
#include <laspack/rtc.h>
#include <xc/getopts.h>
}

namespace GW {
	
typedef Vector LSP_Vector;
typedef QMatrix LSP_Matrix;

/*------------------------------------------------------------------------------*/
/** 
 *  \class  GW_SparseMatrix
 *  \brief  A sparse matrix.
 *  \author Gabriel Peyr
 *  \date   6-8-2003
 *
 *  It is a wrapper \c LASPACK library.
 */ 
/*------------------------------------------------------------------------------*/

std::ostream& operator<<(std::ostream &s, class GW_SparseMatrix& A);

class GW_SparseMatrix
{
public:

	enum T_IterativeSolverType
	{	
		/* classical iterative methods */
		IterativeSolver_Jacobi,
		IterativeSolver_SORForw,
		IterativeSolver_SORBackw,
		IterativeSolver_SSOR,
		/* semi-iterative methods */
		IterativeSolver_Chebyshev,
		/* CG and CG-like methods */
		IterativeSolver_CG,
		IterativeSolver_CGN,
		IterativeSolver_GMRES,
		IterativeSolver_BiCG,
		IterativeSolver_QMR,
		IterativeSolver_CGS,
		IterativeSolver_BiCGSTAB,
	};

	enum T_PreconditionerType
	{
		Preconditioner_Jacobi,
		Preconditioner_SSOR,
		Preconditioner_ILU,
		Preconditioner_NULL			// no preconditioner
	};

	GW_SparseMatrix( GW_U32 nSize, ElOrderType ElOrder = Rowws )
	{
		Q_Constr( &M_,  "GW_SparseMatrix", nSize, False, ElOrder, Normal, True );
		RowPosition_ = new GW_U32[nSize];
		memset( RowPosition_, 0, nSize*sizeof(GW_U32) );
	}
	GW_SparseMatrix( GW_SparseMatrix& Original )
	{
		Q_Constr( &M_,  "GW_SparseMatrix", Original.GetDim(), False, Q_GetElOrder( &Original.M_ ), Normal, True );
		RowPosition_ = new GW_U32[Original.GetDim()];
		memset( RowPosition_, 0, Original.GetDim()*sizeof(GW_U32) );
		(*this) = Original;
	}
	GW_SparseMatrix( GW_MatrixNxP& Original, ElOrderType ElOrder = Rowws  )
	{
		GW_ASSERT( Original.GetNbrRows()==Original.GetNbrCols() );
		GW_U32 nSize = GW_MIN(Original.GetNbrRows(),Original.GetNbrCols());
		Q_Constr( &M_,  "GW_SparseMatrix", nSize, False, ElOrder, Normal, True );
		RowPosition_ = new GW_U32[nSize];
		memset( RowPosition_, 0, nSize*sizeof(GW_U32) );
		/* copy data */
		this->BuildFromFull(Original);
	}
	virtual ~GW_SparseMatrix()
	{
		Q_Destr( &M_ );
		GW_DELETEARRAY(RowPosition_);
	}
	GW_SparseMatrix& operator =( GW_SparseMatrix& Original )
	{
		GW_ASSERT( this->GetDim()==Original.GetDim() );
		GW_SparseMatrix::Copy( *this, Original.M_ );
		return *this;
	}
	
	void BuildFromFull( GW_MatrixNxP& Original )
	{
		GW_U32 nSize = Original.GetNbrRows();
		/* copy data */
		for( GW_U32 i=0; i<nSize; ++i )
		{
			/* count number of non-zero elements */
			GW_U32 nNonZero = 0;
			for( GW_U32 j=0; j<nSize; ++j )
				if( Original.GetData(i,j)!=0.0 )
					nNonZero++;
			this->SetRowSize(i,nNonZero);
			for( GW_U32 j=0; j<nSize; ++j )
				if( Original.GetData(i,j)!=0.0 )
					this->SetData( i,j, Original.GetData(i,j) );
		}
	}

	GW_U32 GetDim()
	{
		return (GW_U32) Q_GetDim(&M_);
	}

	//-------------------------------------------------------------------------
	/** \name accessors */
	//-------------------------------------------------------------------------
	//@{
	GW_Float GetData(GW_U32 i, GW_U32 j)
	{
		GW_ASSERT( i<this->GetDim() && j<this->GetDim() );
		return Q_GetEl(&M_, i+1,j+1);
	}
	GW_U32 GetRowSize( GW_U32 i )
	{
		return (GW_U32) Q_GetLen( &M_, i+1 );
	}
	void SetRowSize( GW_U32 i, GW_U32 nSize )
	{
		GW_ASSERT( i<this->GetDim() );
		GW_ASSERT( nSize<=this->GetDim() );
		Q_SetLen( &M_, i+1, nSize );
		// restart the counter ...
		RowPosition_[i] = 0;
	}
	void ResetRowSize( GW_U32 i, GW_U32 nSize )
	{
		GW_ASSERT( i<this->GetDim() );
		GW_ASSERT( nSize<=this->GetDim() );
		/* we must save the previous stuff */
		std::list<GW_U32> PosList;
		std::list<GW_Float> ValList;
		for( GW_U32 k=0; k<RowPosition_[i]-1; ++k )
		{
			ValList.push_back( Q_GetVal( &M_, i+1, k ) );
			PosList.push_back( (GW_U32) Q_GetPos( &M_, i+1, k ) );
		}
		/* now reset the length */
		Q_SetLen( &M_, i+1, nSize );
		RowPosition_[i] = 0;
		/* now reset the previous entry */
		std::list<GW_Float>::iterator it_val = ValList.begin();
		for( std::list<GW_U32>::iterator it_pos = PosList.begin(); it_pos!=PosList.end(); ++it_pos )
		{
			GW_ASSERT( it_val!=ValList.end() );
			GW_Float val = *it_val;
			GW_U32 pos = *it_pos;	GW_ASSERT( pos>0 );
			this->SetData( i, pos-1, val );
			++it_val;
		}
		GW_ASSERT( it_val==ValList.end() );
	}
	void SetData( GW_U32 i, GW_U32 j, GW_Float rVal )
	{
		GW_ASSERT( i<this->GetDim() && j<this->GetDim() );
		GW_ASSERT( this->GetRowSize(i)>0 && RowPosition_[i]<this->GetRowSize(i) );
		RowPosition_[i]++;
		Q_SetEntry(&M_, i+1, RowPosition_[i]-1, j+1, rVal);	// matrix, row, entry, col, val
	}
	//@}

	/** access element by element */
	GW_Float AccessEntry( GW_U32 i, GW_U32 entry, GW_U32& j )
	{
		GW_ASSERT( entry<this->GetRowSize(i) );
		if( entry<this->GetRowSize(i) )
		{
			j = (GW_U32) Q_GetPos(&M_, i+1, entry) - 1;
			return Q_GetVal(&M_, i+1, entry);
		}
		return 0;
	}

	//-------------------------------------------------------------------------
	/** \name multiplication methods (both optimised and non-optimised) */
	//-------------------------------------------------------------------------
	//@{
	GW_VectorND operator*(GW_VectorND& v)
	{
		GW_VectorND r(v.GetDim());
		GW_SparseMatrix::Multiply( *this, v, r );
		return r;
	}
	static void Multiply(GW_SparseMatrix& a, GW_VectorND& v, GW_VectorND& r)
	{
		GW_ASSERT( a.GetDim()==v.GetDim() );
		GW_ASSERT( a.GetDim()==r.GetDim() );
		LSP_Vector v1;
		V_Constr( &v1, "temp", v.GetDim(), Normal, True );
		GW_SparseMatrix::Copy( v1, v );
		Asgn_VV( &v1, Mul_QV( &a.M_, &v1 ) );
		GW_SparseMatrix::Copy( r, v1 );
		V_Destr( &v1 );
	}
	//@}

	GW_SparseMatrix Transpose()
	{
		ElOrderType ElOrder = Q_GetElOrder(&M_);
		if( ElOrder == Rowws )
            ElOrder = Clmws;
        else if( ElOrder == Clmws )
            ElOrder = Rowws;
		GW_SparseMatrix r(this->GetDim(), ElOrder);
		GW_SparseMatrix::Transpose( *this, r );
		return r;
	}
	static void  Transpose(GW_SparseMatrix& a, GW_SparseMatrix& r)
	{
		GW_SparseMatrix::Copy( r, *Transp_Q(&a.M_) );
	}

	void BuildTriDiag( GW_Float val_down, GW_Float val_up, GW_Float val_diag )
	{
		GW_U32 n = this->GetDim();
		this->SetRowSize(0,2);
		this->SetData(0,0, val_diag);
		this->SetData(0,1, val_up);
		for( GW_U32 i=1; i<n-1; ++i )
		{
			this->SetRowSize(i,3);
			this->SetData(i,i-1, val_down);
			this->SetData(i,i, val_diag);
			this->SetData(i,i+1, val_up);
		}
		this->SetRowSize(n-1,2);
		this->SetData(n-1,n-2, val_down);
		this->SetData(n-1,n-1, val_diag);
	}

	static void TestClass(std::ostream &s = cout)
	{
		TestClassHeader("GW_SparseMatrix", s);
		GW_U32 n = 10;

		s << "Creation test." << endl;
		GW_SparseMatrix SM(n);
		SM.BuildTriDiag( 3,3,4 );
		s << SM;

		s << "M1=M2 : Copy test." << endl;
		GW_SparseMatrix SM2(n);
		SM2 = SM;
		s << SM;
		
		s << "M^T : Transpose test." << endl;
		s << SM.Transpose();

		
		s << "M*[1] : Matrix/Vector multiplication test." << endl;
		GW_VectorND v(n);
		v.SetValue(1);
		s << SM*v << endl;

		n = 1000;
		GW_Float eps = 1e-10;
		GW_SparseMatrix SM3(n);
		SM3.BuildTriDiag( 3,3,4 );
		s << "Conjugate gradient with SSOR test, size " << n << endl;
		GW_VectorND x(n);
		GW_VectorND b(n);
		b.Randomize();
		GW_Float err = SM3.IterativeSolve( x, b, GW_SparseMatrix::IterativeSolver_CG, GW_SparseMatrix::Preconditioner_SSOR, 
				eps, 5*n, 1.2, GW_True );
		s << "Error=" << err << "." << endl;

		
		GW_SparseMatrix SM4(n);
		SM4.BuildTriDiag( -3,3,4 );
		s << "Bi-conjugate gradient with SSOR test, size " << n << endl;
		b.Randomize();
		err = SM4.IterativeSolve( x, b, GW_SparseMatrix::IterativeSolver_BiCG, GW_SparseMatrix::Preconditioner_SSOR, 
				eps, 5*n, 1.2, GW_True );
		s << "Error=" << err << "." << endl;
		TestClassFooter("GW_SparseMatrix", s);
	}

	void Print(std::ostream &s)
	{
		s << "Sparse matrix, size " << this->GetDim() << "x" << this->GetDim() << "." << endl; 
		for( GW_U32 i=0; i<this->GetDim(); ++i )
		{
			s << "|";
			for( GW_U32 j=0; j<this->GetDim();++j )
			{
				if( Q_GetElOrder(&M_)==Rowws )
					s << this->GetData(i,j);
				else
					s << this->GetData(j,i);
				if( j!=this->GetDim()-1 )
					s << " ";
			}
			s << "|" << endl;
		}
	}

	GW_Float IterativeSolve( GW_VectorND& x, GW_VectorND& b, T_IterativeSolverType Solver, T_PreconditionerType Preconditioner, 
						 GW_Float eps = 1e-10, GW_U32 nMaxIter = 100, GW_Float Omega = 1.2, GW_Bool bGetError = GW_False )
	{
		LSP_Vector xv;
		LSP_Vector bv;
		size_t Dim = this->GetDim();

		/* temportary data */
		V_Constr( &xv, "xv", Dim, Normal, True );
		GW_SparseMatrix::Copy( xv, x );
		V_Constr( &bv, "bv", Dim, Normal, True );
		GW_SparseMatrix::Copy( bv, b );

		/* set the preconditioner */
		PrecondProcType PrecondProc = NULL;
		switch( Preconditioner ) {
		case Preconditioner_SSOR:
			PrecondProc = SSORPrecond;
			break;
		case Preconditioner_Jacobi:
			PrecondProc = JacobiPrecond;
			break;
		case Preconditioner_ILU:
			PrecondProc = ILUPrecond;
			break;
		}

		SetRTCAccuracy( eps );
		V_SetAllCmp( &xv, 0 );	// initial guess
		switch(Solver) {
		case IterativeSolver_Jacobi:
			JacobiIter(&M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_SORForw:
			SORForwIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_SORBackw:
			SORBackwIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_SSOR:
			SSORIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_Chebyshev:
			ChebyshevIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_CG:
			CGIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_CGN:
			CGNIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_GMRES:
			GMRESIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_BiCG:
			BiCGIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_QMR:
			QMRIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_CGS:
			CGSIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		case IterativeSolver_BiCGSTAB:
			BiCGSTABIter( &M_, &xv, &bv, nMaxIter, PrecondProc, Omega );
			break;
		default:
			GW_ASSERT( GW_False );
		}

		GW_SparseMatrix::Copy( x, xv );

		if( bGetError )
		{
			/* r = b - A * x */
			return (b-(*this)*x).Norm2();
		}
		else 
			return 0;

		/* destroy data */
		V_Destr(&xv);
		V_Destr(&bv);
	}

	/** Internal helper */
	static void Copy( GW_VectorND& v1, LSP_Vector& v2 )
	{
		GW_ASSERT( V_GetDim(&v2)==v1.GetDim() );
		for( GW_U32 i=0; i<v1.GetDim(); ++i )
			v1.SetData( i, V_GetCmp( &v2, i+1 ) );
	}
	static void Copy( LSP_Vector& v1, GW_VectorND& v2 )
	{
		GW_ASSERT( V_GetDim(&v1)==v2.GetDim() );
		for( GW_U32 i=0; i<v2.GetDim(); ++i )
			V_SetCmp( &v1, i+1, v2.GetData(i) );
	}
	static void Copy( GW_SparseMatrix& m1, LSP_Matrix& m2 )
	{
		GW_ASSERT( m1.GetDim() == Q_GetDim(&m2) );	
		for( GW_U32 i=1; i<=Q_GetDim(&m2); ++i )
		{
			size_t l = Q_GetLen(&m2, i);
			Q_SetLen( &(m1.M_), i, l );
			for( GW_U32 j=0; j<l; ++j )
			{
				GW_U32 pos = (GW_U32) Q_GetPos( &m2, i, j );
				Q_SetEntry( &m1.M_, i, j, pos, Q_GetVal(&m2, i,j) );
			}
		}
	}

// private:

	LSP_Matrix M_;
	GW_U32* RowPosition_;

};

inline
std::ostream& operator<<(std::ostream &s, GW_SparseMatrix &A)
{
	A.Print(s);

	return s;
}

} // End namespace GW


#endif // _GW_SPARSEMATRIX_H_


///////////////////////////////////////////////////////////////////////////////
//  Copyright (c) Gabriel Peyr
///////////////////////////////////////////////////////////////////////////////
//                               END OF FILE                                 //
///////////////////////////////////////////////////////////////////////////////
