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

#include "precondnl.h"

#include "int.h"
#include "nonlin.h"

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

#include "cmdpars.h"
extern CmdPars Cmd;

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

NonLinearSGGS:: NonLinearSGGS(NonLinearity& nonLinearity0) 

    : Preconditioner(), nonLinearity(nonLinearity0), sol(0)
{ }
//-------------------------------------------------------------------------


void NonLinearSGGS:: initialize(SystemMatrix* A, Vector<Num>& x, Vector<Num>& /*b*/)
{
    int i;
    if (infoLinSystem)  cout << " -- Nonlinear SL-GS Preconditioner\n";

    if (A->DirectSolution()) 
    {
	cout << "\n*** NonLinearSGGS::initialize(...) : wrong Matrix-type;"
	     << "\n\t DirectSolution()==true (set directSolverNodeLimit=0 ?)\n";
	cout.flush(); abort();
    }

    FORALL(x,i) { if (dirichletBCs->isSet(i)) x[i] = dirichletBCs->value(i); }
    sol = &x;

    A->symmetricExpansion();
}
//-------------------------------------------------------------------------

void NonLinearSGGS:: close(SystemMatrix* A, Vector<Num>& /*x*/, Vector<Num>& /*b*/)
{
    A->removeSymmetricExpansion();
}
//-------------------------------------------------------------------------


// -- 	    perform one non-linear Gauss-Seidel step on A*e = r :


void NonLinearSGGS:: invert(Vector<Num>& e, SystemMatrix* A, Vector<Num>& r)
{
    int i;
    Num res;

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

    for (i=1; i<=A->Dim(); ++i)
    {
	if (dirichletBCs->isSet(i)) continue;

	res = A->MultRow(i,e);
	res = r[i] - res;            // the new (linear) residual of node i
	
	nonLinearity.defectCorrection(i, e[i], res, (*sol)[i], A->Diag(i));

	r[i] -= res;			// res may have changed
    }
}
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------


// --	 Non-linear Gauss-Seidel multi-level preconditioner


NonLinearMLGS:: NonLinearMLGS(NonLinearity& nonLinearity0) 

    : MMGPreconditioner(), 
      nonLinearity(nonLinearity0),
      lowDefObstacle(0,DefaultStackSize), uppDefObstacle(0,DefaultStackSize),
      sol(0)
{ 
    topLevelSmoothing = True;	Cmd.get("topLevelSmoothing",&topLevelSmoothing); 

    lowDefObstacle.push(new Vector<Real>(1));
    uppDefObstacle.push(new Vector<Real>(1));
}
//-------------------------------------------------------------------------

NonLinearMLGS:: ~NonLinearMLGS()
{
    int i;
    FORALL(lowDefObstacle,i)
    {
	delete lowDefObstacle[i];
	delete uppDefObstacle[i];
    }
}
//-------------------------------------------------------------------------


void NonLinearMLGS:: update(SystemMatrix* A, FamilyTree* familyTree0,
			    DirichletBCs* dirichletBCs0)
{
    int level, dim;

    familyTree   = familyTree0;
    dirichletBCs = dirichletBCs0;
    initParameters();

    if (A->DirectSolution()) 
    {
	cout << "\n*** NonLinearMLGS::update : wrong Matrix-type;"
	     << "\n\t DirectSolution()==True (set directSolverNodeLimit=0 ?)\n";
	cout.flush(); abort();
    }

    if (maxLevel < 0)  				// level 0 call
    { 
	maxLevel = 0; 
	Al.push(A); 
	lowDefObstacle[0]->resize(A->Dim());
	uppDefObstacle[0]->resize(A->Dim());
	return;
    }	

    maxLevel = familyTree->MaxLevel();

    if (maxLevel > 0)
	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);
    }

    if (maxLevel > uppDefObstacle.high())		// depth has increased
    {
	lowDefObstacle.push(new Vector<Real>(1));
	uppDefObstacle.push(new Vector<Real>(1));
    }
    for (level=0; level<=maxLevel; ++level)
    {
	dim = Al[level]->Dim();
	lowDefObstacle[level]->resize(dim);
	uppDefObstacle[level]->resize(dim);
    }

    if (infoPrecond)  memSpace(True);
    if (Cmd.isTrue("printAL"))  { int i; FORALL(Al,i) Al[i]->print(); }
}
//-------------------------------------------------------------------------


void NonLinearMLGS:: initialize(SystemMatrix* A, Vector<Num>& x, Vector<Num>& /*b*/)
{
    int i;
    if (infoLinSystem)  cout << " -- Nonlinear ML-GS Preconditioner\n";

    aux.resize(A->Dim());
    eMG.resize(A->Dim());
    ADiag0.resize(A->Dim());

    FORALL(x,i) { if (dirichletBCs->isSet(i)) x[i] = dirichletBCs->value(i); }
    sol = &x;

    A->symmetricExpansion();
}
//-------------------------------------------------------------------------

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


void NonLinearMLGS:: invert(Vector<Num>& e, SystemMatrix* A, Vector<Num>& r)
{
    int i;

    FORALL(ADiag0,i) ADiag0[i] = A->Diag(i);		// save diagonal
    
    NonLinGSStep(e, A, r);

    residual(aux,r,*A,e);				// new residual in aux
    FORALL(r,i) r[i] = aux[i];

    el.push(&eMG);     			// available for restrictObstacles

    MGCycle(eMG, *Al[maxLevel], r, maxLevel);

    el.pop();

    FORALL(e,i) e[i] += eMG[i];		     // add multi-grid defect correction

    FORALL(ADiag0,i) A->Diag(i) = ADiag0[i];
}
//-------------------------------------------------------------------------


void NonLinearMLGS:: NonLinGSStep(Vector<Num>& e, SystemMatrix* A, 
				  Vector<Num>& r)
{
    int  i;
    Num  res, a;
    Bool critic;
    Vector<Real>& uppDefO = *uppDefObstacle[maxLevel];
    Vector<Real>& lowDefO = *lowDefObstacle[maxLevel];

    FORALL(e,i)  
    {
	e[i] = uppDefO[i] = lowDefO[i] = 0.0;
	setUnCritical(i);
    }

    for (i=1; i<=A->Dim(); ++i)
    {
	if (dirichletBCs->isSet(i)) continue;

	res = A->MultRow(i,e);
	res = r[i] - res;          // the new (linear) residual of node i
	
	nonLinearity.defectCorrection(i, e[i], res, (*sol)[i], A->Diag(i), 
				      a, uppDefO[i], lowDefO[i], critic);
	r[i] -= res;
	A->Diag(i) += a;
	
	if (critic) setCritical(i); 
    }
}
//-------------------------------------------------------------------------


void NonLinearMLGS:: preSmooth(int level, Vector<Num>& e, SystemMatrix& A, 
			       Vector<Num>& r)
{
    if (level==maxLevel && !topLevelSmoothing) return;

    int i, n;
    Vector<Real>& uppO = *uppDefObstacle[level];
    Vector<Real>& lowO = *lowDefObstacle[level];

    for (n=1; n<=nPreSmooth; ++n)
    {
	for (i=1; i<=A.Dim(); ++i)
	{
	    if (!isCritical(i,level))
	    {
		A.smoothNode(i,e,r);
		
		if (e[i] > uppO[i])  e[i] = uppO[i]; 
		if (e[i] < lowO[i])  e[i] = lowO[i];
	    }
	}
    }
}
//-------------------------------------------------------------------------


void NonLinearMLGS:: postSmooth(int level, Vector<Num>& e, SystemMatrix& A, 
				Vector<Num>& r)
{
    if (level==maxLevel && !topLevelSmoothing) return;

    int i, n;
    Vector<Real>& uppO = *uppDefObstacle[level];
    Vector<Real>& lowO = *lowDefObstacle[level];

    for (n=1; n<=nPostSmooth; ++n)
    {
	for (i=A.Dim(); i>=1; --i)
	{
	    if (!isCritical(i,level))
	    {
		A.smoothNode(i,e,r);
		
		if (e[i] > uppO[i])  e[i] = uppO[i]; 
		if (e[i] < lowO[i])  e[i] = lowO[i];  
	    }
	}
    }
}
//-------------------------------------------------------------------------


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

    restrictObstacles(level);

    MLMatrix* AMLM1 = Al[level-1]->castToMLMatrix();
    MLMatrix* AML   = Al[level]->castToMLMatrix();

    AMLM1->markNodes();			// -> all nodes will be restricted
    AMLM1->GalerkinRestriction(*AML, *familyTree->getGeneration(level));
    AMLM1->symmetricExpansion();
}
//-------------------------------------------------------------------------


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


// --	 Truncated non-linear Gauss-Seidel multi-level preconditioner


TrcNonLinearMLGS:: TrcNonLinearMLGS(NonLinearity& nonLinearity0) 

    : NonLinearMLGS(nonLinearity0), critical(0,DefaultStackSize)
{ 
    critical.push(new Vector<SBool>(1));
}
//-------------------------------------------------------------------------

TrcNonLinearMLGS:: ~TrcNonLinearMLGS()
{
    int i;
    FORALL(critical,i) delete critical[i];
}
//-------------------------------------------------------------------------


void TrcNonLinearMLGS:: update(SystemMatrix* A, FamilyTree* familyTree0,
			       DirichletBCs* dirichletBCs0)
{
    int level;

    NonLinearMLGS::update(A, familyTree0, dirichletBCs0);

    if (maxLevel > critical.high())		// depth has increased
      critical.push(new Vector<SBool>(1));

    for (level=0; level<=maxLevel; ++level)
    {
	Vector<SBool>& critic = *critical[level];
	critic.resize(Al[level]->Dim());
    }
}
//-------------------------------------------------------------------------


void TrcNonLinearMLGS:: initialize(SystemMatrix* A, Vector<Num>& x, 
				   Vector<Num>& /*b*/)
{
    int i;
    if (infoLinSystem)  cout << " -- Truncated Nonlinear ML-GS Prec.\n";

    aux.resize(A->Dim());
    eMG.resize(A->Dim());
    ADiag0.resize(A->Dim());

    FORALL(x,i) { if (dirichletBCs->isSet(i)) x[i] = dirichletBCs->value(i); }
    sol = &x;
    
    A->symmetricExpansion();
}
//-------------------------------------------------------------------------

void TrcNonLinearMLGS:: restr(Vector<Num>& rHigh, Vector<Num>& rLow, 
				 int level)
{
    int i;
    Son* son;
    Generation& generation = *familyTree->getGeneration(level);
    Vector<SBool>& criticM1 = *critical[level-1];
    Vector<SBool>& critic   = *critical[level];

    FORALL(criticM1,i) criticM1[i] = 0;
    FORALL(critic,i) 
    {
	son = generation.getSon(i);
	if (son->NoOfFathers() == 1)  criticM1[son->getFather()] = critic[i];
    }

    FORALL(critic,i) { if (critic[i])  rHigh[i] = 0.0; }

    familyTree->restr(rHigh, rLow, level);
    Al[level-1]->resetDirichletEntries(rLow);

    restrictObstacles(level);
    
    MLMatrix* AMLM1 = Al[level-1]->castToMLMatrix();
    MLMatrix* AML   = Al[level]->castToMLMatrix();

    AMLM1->markNodes();			// -> all nodes will be restricted
    AMLM1->GalerkinRestriction(*AML, generation, &critic);
    AMLM1->symmetricExpansion();
}
//-------------------------------------------------------------------------


void TrcNonLinearMLGS:: prolong(Vector<Num>& eLow, Vector<Num>& eHigh, int level)
{
    int i;
    Vector<SBool>& criticM1 = *critical[level-1];
    Vector<SBool>& critic   = *critical[level];

    FORALL(eLow,i) { if (criticM1[i]) eLow[i] = 0.0; }
    Al[level-1]->resetDirichletEntries(eLow);

    familyTree->prolong(eLow, eHigh, level);

    FORALL(eHigh,i) { if (critic[i]) eHigh[i] = 0.0; }
}
//-------------------------------------------------------------------------

void TrcNonLinearMLGS:: setCritical(int node)  
					{ (*critical[maxLevel])[node] = True; }
void TrcNonLinearMLGS:: setUnCritical(int node) 
					{ (*critical[maxLevel])[node] = False; }
Bool TrcNonLinearMLGS:: isCritical(int node, int level) 
					{ return (*critical[level])[node]; }
//-------------------------------------------------------------------------


void NonLinearMLGS:: restrictObstacles(int highLevel)
{
    int i, father1, father2;
    Real v1,v2,vm, meanv;
    Son* son;
    const Generation* generation = familyTree->getGeneration(highLevel);

    Vector<Real>& uppOM1  = *uppDefObstacle[highLevel-1];
    Vector<Real>& uppO    = *uppDefObstacle[highLevel];
    Vector<Real>& lowOM1  = *lowDefObstacle[highLevel-1];
    Vector<Real>& lowO    = *lowDefObstacle[highLevel];
   
    Vector<Num>& e = *el[highLevel];

    FORALL(e,i) { uppO[i] -= real(e[i]);  lowO[i] -= real(e[i]);  }

    FORALL(uppO,i) 
    {
	son = generation->getSon(i);
	if (son->NoOfFathers() == 1)  
	{
	    father1 = son->getFather();
	    uppOM1[father1] = uppO[i];
	    lowOM1[father1] = lowO[i];
	}
    }
    
    FORALL(uppO,i)
    {
	son = generation->getSon(i);
	if (son->NoOfFathers() > 1)  
	{
	    father1 = son->getFather(1);
	    father2 = son->getFather(2);
	    
	    v1= uppOM1[father1];
	    v2= uppOM1[father2];
	    vm = uppO[i];
	    meanv=(v1+v2)/2.;
            
	    
	    if (meanv != 0.0 && meanv > vm && !isCritical(i,highLevel))
	    {
		if (v1 > vm)
		{
		    if (v2 > vm) 
		    {
			uppOM1[father1] = vm;
			uppOM1[father2] = vm;
		    }
		    else  uppOM1[father1] = 2*vm - v2;
		}
		else
		{
		    if (v2 > vm) uppOM1[father2] = 2*vm - v1;
		}
	    }
	    
	    v1= lowOM1[father1];
	    v2= lowOM1[father2];
	    vm = lowO[i];
	    meanv=(v1+v2)/2.;
	    
	    //if (equal(meanv,0.0)) continue;
	    //if (meanv >= vm)      continue;

	    if (meanv != 0.0 && meanv < vm && !isCritical(i,highLevel))
	    {
		if (v1 < vm)
		{
		    if (v2 < vm) 
		    {
			lowOM1[father1] = vm;
			lowOM1[father2] = vm;
		    }
		    else  lowOM1[father1] = 2*vm - v2;
		}
		else
		{
		    if (v2 < vm) lowOM1[father2] = 2*vm - v1;
		}
	    }
	}
    }

    // --      restore old values for post-smoothing:

    FORALL(e,i) { uppO[i] += real(e[i]);  lowO[i] += real(e[i]);  } 
}