/*
 $Id: precondmg.cc,v 1.4 1996/11/20 09:57:27 roitzsch Exp $
 (C)opyright 1996 by Konrad-Zuse-Center, Berlin
 All rights reserved.
 Part of the Kaskade distribution
*/

#include "precondmg.h"

#include "numerics.h"
#include "utils.h"

#include "sysmat.h"
#include "sysmatml.h"
#include "family.h"
#include "dirichlet.h"

#include "cmdpars.h"
extern CmdPars Cmd;

//-------------------------------------------------------------------------


MGPreconditioner:: MGPreconditioner() :  Preconditioner(), 

	Al(0,DefaultStackSize), rl(0,DefaultStackSize), 
	el(0,DefaultStackSize), omega(0,DefaultStackSize)
{
    maxLevel = -1;

    baseLevel = 0;	Cmd.get("BaseLevel", &baseLevel);
    if (baseLevel < 0) { cout << "\n*** MG: BaseLevel must be >= 0\n"; cout.flush(); abort(); }

    nPreSmooth  = 1;	Cmd.get("nPreSmooth", &nPreSmooth); 
    nPostSmooth = 1;	Cmd.get("nPostSmooth",&nPostSmooth);
    nCycle      = 1;	// Cmd.get("nCycle",     &nCycle);  W-Cycle not allowed

    nodeProgFactor = 2.0;  Cmd.get("MGNodeProgFactor", &nodeProgFactor); 
}
//-------------------------------------------------------------------------

MGPreconditioner:: ~MGPreconditioner()
{
    int i;
    FORALL(Al,i) delete Al[i];
    FORALL(rl,i) delete rl[i];
    FORALL(el,i) delete el[i];
}
//-------------------------------------------------------------------------


int MGPreconditioner:: memSpace(int print)
{
    int i, space = omega.memSpace();

    if (mode() == multiLevel)  space += Al.Top()->memSpaceLU();  // <- Allocator

    FORALL(Al,i) { if (Al[i])  space += Al[i]->memSpace(); }
    FORALL(rl,i) { if (rl[i])  space += rl[i]->memSpace(); }
    FORALL(el,i) { if (el[i])  space += el[i]->memSpace(); }

    if(print)  cout << "\n    Space MG: " 
    		    << Form("%1.6f",float(space)/1e6) << " MB\n";
    return space;
}
//-------------------------------------------------------------------------


void MGPreconditioner:: update(SystemMatrix* A, FamilyTree* familyTree0,
			       DirichletBCs* dirichletBCs0)
{
    updateMG(A, familyTree0, dirichletBCs0);
}
//-------------------------------------------------------------------------


// --	        update for the step-by-step version


void MGPreconditioner:: updateMG(SystemMatrix* A, FamilyTree* familyTree0,
				 DirichletBCs* dirichletBCs0)
{
    int i;

    dirichletBCs = dirichletBCs0;
    familyTree   = familyTree0;

    ++maxLevel;

    initParameters();

    if (maxLevel > 0)
    {
	rl.push(new Vector<Num>(dimM1));
	el.push(new Vector<Num>(dimM1));

	if (diagonalOnly()) if (Al[maxLevel-1])  Al[maxLevel-1]->removeLU();
    }

    skipLevel(A);	// remove previous level matrix due to node progression

    dimM1 = A->Dim();


    if (infoPrecond)  
    {
	memSpace(True);
	if (maxLevel == 0)
	{
	    cout << "    MG: nPre=" << nPreSmooth << " nPost=" << nPostSmooth 
	     	 << " nCycle=" << nCycle << "  baseLevel=" << baseLevel 
	     	 << "  nodeProgFactor=" << nodeProgFactor << "\n";
	}
    }
    if (Cmd.isTrue("printAL"))  
    { 	FORALL(Al,i) { cout << "\nLevel " << i << ": "; Al[i]->print(); } }
}
//-------------------------------------------------------------------------


void MGPreconditioner:: skipLevel(SystemMatrix* A)
{
    if (maxLevel == 0)  
    { 
	Al.push(A); 
	lastDim = A->Dim(); 
	skipLevelM1 = False;
    }
    else
    {
	if (skipLevelM1) 			// level-1 to be skipped, if 
	{					//    matrix is not factorized
	    if (!Al.Top()->DirectSolution())
	    {
		delete Al.pop();
		Al.push(0);
	    }
	}	
	Al.push(A); 
	
	if (A->Dim() > nodeProgFactor*lastDim)  
	{ 
	    lastDim = A->Dim(); 
	    skipLevelM1 = False;
	}
	else  				// this level to be skipped on next run
	{ 
	    skipLevelM1 = True;
	    if (infoPrecond) 
	       cout << "\n    MG: matrix level " << maxLevel << " skipped\n";
	}
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void MGPreconditioner:: updateML(SystemMatrix* A, FamilyTree* familyTree0,
				 DirichletBCs* dirichletBCs0)
{
    int i, level, dim;

    dirichletBCs = dirichletBCs0;
    familyTree   = familyTree0;

    initParameters();

    if (maxLevel < 0)  				// level 0 call
    { 
	Al.push(A); 
	maxLevel = 0;  
	return; 
    }	

    maxLevel = familyTree->MaxLevel();

    if (maxLevel == 0)
    {
	cout <<"\n*** MGPreconditioner:: updateML: error in maxLevel (==0)\n";
	cout.flush(); abort();
    }
    else
    {
	if (maxLevel-1 > rl.high())		// depth has increased
	{
	    Al.push(A);
	    rl.push(new Vector<Num>(1));
	    el.push(new Vector<Num>(1));
	}

	for (level=0; level<maxLevel; ++level)
	{
	    dim = Al[level]->Dim();
	    rl[level]->resize(dim);
	    el[level]->resize(dim);
	}
    }

    GalerkinRestriction();
    

    if (infoPrecond)  
    {
	memSpace(True);
	if (maxLevel == 0)
	{
	    cout << "    MG: nPre=" << nPreSmooth << " nPost=" << nPostSmooth 
	     	 << " nCycle=" << nCycle << "  baseLevel=" << baseLevel  << "\n";
	}
    }
    if (Cmd.isTrue("printAL"))  
    { 	FORALL(Al,i) { cout << "\nLevel " << i << ": "; Al[i]->print(); } }
}
//-------------------------------------------------------------------------


void MGPreconditioner:: GalerkinRestriction()
{
    int level;
    MLMatrix* AML, *AMLP1;

    for (level=maxLevel-1; level >= 0; --level)    
    { 
	if (Al[level]->DirectSolution())  break;

	AML   = Al[level]->castToMLMatrix();
	AMLP1 = Al[level+1]->castToMLMatrix();

	AML->markNodes();	      // only marked nodes will be restricted
	AML->GalerkinRestriction(*AMLP1, *familyTree->getGeneration(level+1));

	if (level < maxLevel-1)
	{
	    if (diagonalOnly()) AMLP1->removeLU();
	    AMLP1->removeNonSmoothedEntries();
	}

	if (level == 0)
	{
	    if (diagonalOnly()) AML->removeLU();
	    AML->removeNonSmoothedEntries();
	}
    }
};
//-------------------------------------------------------------------------


void MGPreconditioner:: invert(Vector<Num>& e, SystemMatrix* /*A*/, Vector<Num>& r)
{
    MGCycle(e, *Al[maxLevel], r, maxLevel);
}
//-------------------------------------------------------------------------


void MGPreconditioner:: cleanStacks(int level)
{
    for (int i=0; i<=level; ++i)
    {
	delete Al[i];  Al[i] = 0; 
	delete el[i];  el[i] = 0; 
	delete rl[i];  rl[i] = 0; 
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void MGPreconditioner:: restr(Vector<Num>& rHigh, Vector<Num>& rLow, 
				 int level)
{
    familyTree->restr(rHigh, rLow, level);
    if (Al[level-1]) Al[level-1]->resetDirichletEntries(rLow);
}
//-------------------------------------------------------------------------


void MGPreconditioner:: prolong(Vector<Num>& eLow, Vector<Num>& eHigh, int level)
{
    if (Al[level-1]) Al[level-1]->resetDirichletEntries(eLow);
    familyTree->prolong(eLow, eHigh, level);
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void MMGPreconditioner:: close(SystemMatrix* /*A*/, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    aux.resize(1);
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


MMGPreconditioner:: MMGPreconditioner() : MGPreconditioner(), aux(1)  { }

//-------------------------------------------------------------------------


void MMGPreconditioner:: MGCycle(Vector<Num>& e, SystemMatrix& A, Vector<Num>& r,
				 int level)
{
    int i, n;

    FORALL(e,i) e[i] = 0.0;			

    if (&A == 0)	  	// no smoothing procedure on this level
    {
	if (level > baseLevel)
	{
	    restr(r, *rl[level-1], level);
	    for (n=1; n<=nCycle; ++n)  MGCycle(*el[level-1], *Al[level-1], 
					       *rl[level-1], level-1);
	    prolong(*el[level-1], e, level);
	}
    }

    else if (A.DirectSolution()) A.FBSubst(e,r); 
    

    else		 // the usual MMG cycle:
    {    
	if (nPreSmooth) preSmooth(level, e, A, r);
		
	if (level > baseLevel)
	{
	    residual(aux, r, A, e);
	    restr(aux, *rl[level-1], level);
	    
	    for (n=1; n<=nCycle; ++n)  MGCycle(*el[level-1], *Al[level-1], 
					       *rl[level-1], level-1);
	    
	    prolong(*el[level-1], e, level);
	}

	if (nPostSmooth) postSmooth(level, e, A, r);
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void AMGPreconditioner:: MGCycle(Vector<Num>& e, SystemMatrix& A, 
				 Vector<Num>& r, int level)
{
    int i;

    FORALL(e,i) e[i] = 0.0;			

    if (&A == 0)	  	// no smoothing procedure on this level
    {
	if (level > baseLevel)
	{
	    restr(r, *rl[level-1], level);
	    MGCycle(*el[level-1], *Al[level-1], *rl[level-1], level-1);
	    prolong(*el[level-1], e, level);
	}
    }

    else if (A.DirectSolution())  A.FBSubst(e,r); 


    else		// the usual AMG cycle:
    {
	smooth(level, e, A, r);

	if (level > baseLevel)
	{
	    restr(r, *rl[level-1], level);
	    MGCycle(*el[level-1], *Al[level-1], *rl[level-1], level-1);
	    prolong(*el[level-1], e, level);
	}
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void MGJacobi:: initParameters()
{
    Real omegaNew;
    if (!Cmd.get("Omega", &omegaNew))  omegaNew = 0.66;
    omega.push(omegaNew);
}

void MGJacobi:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)
	cout << " -- Jacobi MG Precond. (omega = " << omega.Top() <<")\n";

    if (omega.Top() > 0.7) cout <<"\n * omega too big ?\n";

    aux.resize(A->Dim());
}
//-------------------------------------------------------------------------


void MGSGS:: initialize (SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)  cout << " -- SGS MG Preconditioner\n";
    aux.resize(A->Dim());
}
//-------------------------------------------------------------------------


void MGSSOR:: initParameters()
{
    Real omegaNew;
    if (!Cmd.get("Omega",&omegaNew))  omegaNew = 1.0;
      // omegaNew = Al[maxLevel]->omegaOpt(multiGrid);
    omega.push(omegaNew);
}

void MGSSOR:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)
	cout << " -- SSOR MG Precond. (omega = " << omega.Top() <<")\n"; 

    aux.resize(A->Dim());
}
//-------------------------------------------------------------------------


void MGJacobi:: preSmooth(int level, Vector<Num>& e, SystemMatrix& A, 
			  Vector<Num>& r) 
{
    int i, n;
    
    A.DiagDiv(e,r,omega[level]);

    for (n=nPreSmooth-1; n > 0; --n)
    {
	residual(aux,r,A,e);
	A.DiagDiv(aux,aux,omega[level]);	// non-smoothed nodes must
						// be set to zero!
	FORALL(e,i) e[i] += aux[i];
    }
}
//-------------------------------------------------------------------------

void MGJacobi:: postSmooth(int level, Vector<Num>& e, SystemMatrix& A,
			   Vector<Num>& r) 
{
    int i, n;

    for (n=nPostSmooth; n > 0; --n)
    {
	residual(aux,r,A,e);
	A.DiagDiv(aux,aux,omega[level]);	// non-smoothed nodes must
						// be set to zero!
	FORALL(e,i) e[i] += aux[i];
    }
}
//-------------------------------------------------------------------------


void MGSGS:: preSmooth(int /*level*/, Vector<Num>& e, SystemMatrix& A,
		       Vector<Num>& r)
{
    int i, n;
    
    A.Fm1(e,r);
    for (n=nPreSmooth-1; n > 0; --n)
    {
	A.UMult(aux,e);
	FORALL(r,i) aux[i] = r[i]-aux[i]; 	
	A.Fm1(e,aux);
    }
}
//-------------------------------------------------------------------------


void MGSGS:: postSmooth(int /*level*/, Vector<Num>& e, SystemMatrix& A, 
			Vector<Num>& r)
{
    int i, n;

    for (n=nPostSmooth; n > 0; --n)
    {
	A.LMult(aux,e);
	FORALL(r,i) aux[i] = r[i]-aux[i];
	A.FmT(e,aux);
    }
}
//-------------------------------------------------------------------------


void MGSSOR:: preSmooth(int level, Vector<Num>& e, SystemMatrix& A, 
			Vector<Num>& r)
{
    int i, n;
    Real oneMinusOmega = 1.0-omega[level];

    A.Fm1(e,r,omega[level]);
    if (!equal(omega[level],1.0))  { FORALL(e,i) e[i] *= omega[level]; }
	
    for (n=nPreSmooth-1; n > 0; --n)
    {
	A.UMult(aux,e);
	A.DiagMult(e,e,oneMinusOmega);
	FORALL(e,i) aux[i] = e[i] + omega[level] * (r[i] - aux[i]);
	A.Fm1(e,aux,omega[level]);
    }
}
//-------------------------------------------------------------------------


void MGSSOR:: postSmooth(int level, Vector<Num>& e, SystemMatrix& A, 
			 Vector<Num>& r)
{
    int i, n;
    Real oneMinusOmega = 1.0-omega[level];

    for (n=nPostSmooth; n > 0; --n)
    {
	A.LMult(aux,e);
	A.DiagMult(e,e,oneMinusOmega);
	FORALL(e,i) aux[i] = e[i] + omega[level] * (r[i] - aux[i]);
	A.FmT(e,aux,omega[level]);
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void MGCG:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)  cout << " -- CG MG Preconditioner \n";
    p.resize(A->Dim());
    r.resize(A->Dim());
    aux.resize(A->Dim());
}
//-------------------------------------------------------------------------

void MGCG:: close(SystemMatrix* /*A*/, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    p.resize(1);
    r.resize(1);
    aux.resize(1);
}
//-------------------------------------------------------------------------

void MGCG:: preSmooth(int /*level*/, Vector<Num>& x, SystemMatrix& A, Vector<Num>& b)
{
    int i;
    const int h0 = aux.h;
    aux.h = p.h = r.h = A.Dim();
    
    FORALL(r,i) r[i] = b[i];
    A.DiagDiv(aux,r);
    FORALL(p,i) p[i] = aux[i];

    CGIterate(x, A, nPreSmooth);

    aux.h = p.h = r.h = h0;
}
//-------------------------------------------------------------------------

void MGCG:: postSmooth(int /*level*/, Vector<Num>& x, SystemMatrix& A,Vector<Num>& b)
{
    int i;
    const int h0 = aux.h;
    aux.h = p.h = r.h = A.Dim();
    
    A.Mult(aux,x);
    A.resetDirichletEntries(aux);

    FORALL(r,i) r[i] = b[i]-aux[i];
    A.DiagDiv(aux,r);
    FORALL(p,i) p[i] = aux[i];

    CGIterate (x, A, nPostSmooth);

    aux.h = p.h = r.h = h0;
}
//-------------------------------------------------------------------------


void MGCG:: CGIterate(Vector<Num>& x, SystemMatrix& A, int nSmooth)
{
    int iter, i;
    Num alpha, beta, rr, rrM1;

    rr = dot(aux,r);

    for (iter=1; 1; ++iter)
    {
	A.Mult(aux,p);
	A.resetDirichletEntries(aux);
	alpha = rr / dot(p,aux);

	FORALL(x,i) x[i] += alpha * p[i];
	
	if (iter==nSmooth) break;
	
	FORALL(r,i) r[i] -= alpha * aux[i];
	A.DiagDiv(aux,r);

	rrM1 = rr;
	rr   = dot(aux,r);
	beta = rr / rrM1;
	FORALL(p,i) p[i] = aux[i] + beta*p[i];
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void AMGJacobi:: initParameters()
{
    Real omegaNew;
    if (!Cmd.get("Omega", &omegaNew))  omegaNew = 1.0;
    omega.push(omegaNew);
}

void AMGJacobi:: initialize(SystemMatrix* /*A*/, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)
         cout << " -- Jacobi AMG Precond. (omega = " << omega.Top() <<")\n";
}
//-------------------------------------------------------------------------

void AMGSGS:: initialize(SystemMatrix* /*A*/, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)  cout << " -- SGS AMG Preconditioner\n";
}
//-------------------------------------------------------------------------


void AMGJacobi:: smooth(int level, Vector<Num>& e, SystemMatrix& A, 
			Vector<Num>& r) 
{
    A.DiagDiv(e,r,omega[level]);
}
//-------------------------------------------------------------------------


void AMGSGS:: smooth(int /*level*/, Vector<Num>& e, SystemMatrix& A, 
		     Vector<Num>& r)
{
    A.Fm1(e,r);
    A.DiagMult(e,e);
    A.FmT(e,e);
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


// --                   Multi-Level preconditioners:


void MLJacobi:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)
	cout << " -- Jacobi ML Precond. (omega = " << omega.Top() <<")\n";

    if (omega.Top() > 0.7) cout <<"\n * omega too big ?\n";

    aux.resize(A->Dim());
    // GalerkinRestriction();
}
//-------------------------------------------------------------------------

void MLSGS:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)  cout << " -- SGS ML Preconditioner\n";
    aux.resize(A->Dim());
    // GalerkinRestriction();
}
//-------------------------------------------------------------------------


void MLSSOR:: initParameters()
{
    Real omegaNew;
    if (!Cmd.get("Omega", &omegaNew))  omegaNew = 1.0;
    omega.push(omegaNew);
}

void MLSSOR:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)
	cout << " -- SSOR ML Precond. (omega = " << omega.Top() <<")\n"; 

    aux.resize(A->Dim());
    // GalerkinRestriction();
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


void MLCG:: initialize(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)  cout << " -- CG ML Preconditioner \n";
    p.resize(A->Dim());
    r.resize(A->Dim());
    aux.resize(A->Dim());
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void AMLJacobi:: initialize(SystemMatrix* /*A*/, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)
         cout << " -- Jacobi AML Precond. (omega = " << omega.Top() <<")\n";
    // GalerkinRestriction();
}

void AMLJacobi:: initParameters()
{
    Real omegaNew;
    if (!Cmd.get("Omega", &omegaNew))  omegaNew = 1.0;
    omega.push(omegaNew);
}
//-------------------------------------------------------------------------

void AMLSGS:: initialize(SystemMatrix* /*A*/, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    if (infoLinSystem)  cout << " -- SGS AML Preconditioner\n";
    // GalerkinRestriction();
}
//-------------------------------------------------------------------------