/*
 * Copyright (C) 2008-2012 Delft University of Technology - All rights reserved.
 *
 * Author: Reijer Idema
 */

/*
 * NonLinearSolver class source.
 */
#include "Preconditioner.h"			// HASA Preconditioner (owned)
#include "ForcingTerms.h"			// HASA ForcingTerms (owned)
#include "NonLinearProblem.h"		// HASA NonLinearProblem (not owned)
#include "NonLinearSolver.h"


//#include <iostream>


NONLINEARSOLVER_NAMESPACE_BEGIN


// Initialize static member objects
const std::string NonLinearSolver::msPetscOptionPrecoType = "-pc_type";
const std::string NonLinearSolver::msPetscOptionPrecoMatrixType = "-pc_mat_type";
const std::string NonLinearSolver::msPetscOptionForcingTermsType = "-ft_type";


/////////////////////////////// PUBLIC ///////////////////////////////////////


//============================= LIFECYCLE ====================================


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::NonLinearSolver"
NonLinearSolver::NonLinearSolver(void)
	: mErr(0), mPetscSnes(PETSC_NULL), mpProblem(NULL),
	  mpPrecoShell(NULL), mpPreco(NULL), mpForcing(NULL),
	  mpOwnedPrecoShell(NULL), mpOwnedPreco(NULL), mpOwnedForcing(NULL),
	  mKspInitialRelativeTolerance(0), mKspInitialAbsoluteTolerance(0),
	  mKspInitialDivergenceTolerance(0), mKspInitialMaxIterationCount(0),
	  mConvergedReason(SNES_CONVERGED_ITERATING), mFunctionNorm(0),
	  mNonLinearIterationCount(0), mLinearIterationCount(0), mFunctionEvaluationCount(0),
	  mIsInitialized(false)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Creating NonLinearSolver object.\n"); CHKERRV(mErr);
	}
} // NonLinearSolver::NonLinearSolver


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::~NonLinearSolver"
NonLinearSolver::~NonLinearSolver(void)
{
	Clear();
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Destroying NonLinearSolver object.\n"); CHKERRV(mErr);
	}
} // NonLinearSolver::~NonLinearSolver


//============================= OPERATIONS ===================================


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::Clear"
void NonLinearSolver::Clear(void)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Clearing NonLinearSolver object data.\n"); CHKERRV(mErr);
		if (mPetscSnes != PETSC_NULL) { mErr = SNESDestroy(&mPetscSnes); mPetscSnes = PETSC_NULL; CHKERRV(mErr); }
	}

	mIsInitialized = false;

	mConvergedReason = SNES_CONVERGED_ITERATING;
	mFunctionNorm = 0;
	mNonLinearIterationCount = 0;
	mLinearIterationCount = 0;
	mFunctionEvaluationCount = 0;

	mpProblem = NULL;
	mpPrecoShell = NULL;
	mpPreco = NULL;
	mpForcing = NULL;

	delete mpOwnedPrecoShell;
	mpOwnedPrecoShell = NULL;
	delete mpOwnedPreco;
	mpOwnedPreco = NULL;
	delete mpOwnedForcing;
	mpOwnedForcing = NULL;

	mKspInitialRelativeTolerance = 0;
	mKspInitialAbsoluteTolerance = 0;
	mKspInitialDivergenceTolerance = 0;
	mKspInitialMaxIterationCount = 0;

	mErr = 0;
} // NonLinearSolver::Clear


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::Initialize"
bool NonLinearSolver::Initialize(void)
{
	if (!PetscUtils::IsInitialized()) return false;
	mErr = PetscInfo(PETSC_NULL, "Initializing NonLinearSolver object data.\n"); CHKERRB(mErr);

	// Clear problem data pointer
	mpProblem = NULL;

	// Set up SNES
	mErr = PetscInfo(PETSC_NULL, "Setting up PETSc SNES object.\n"); CHKERRB(mErr);
	if (mPetscSnes != PETSC_NULL) { mErr = SNESDestroy(&mPetscSnes); mPetscSnes = PETSC_NULL; CHKERRB(mErr); }
	mErr = SNESCreate(PETSC_COMM_WORLD, &mPetscSnes); CHKERRB(mErr);
	mErr = SNESSetTolerances(mPetscSnes, 1.e-6, 0, 0, 100, PETSC_DEFAULT); CHKERRB(mErr);	// Default tolerances from SNESCreate(): 1.e-50, 1.e-8, 1.e-8, 50, 10000

	// Set up KSP and get initial tolerances
	mErr = PetscInfo(PETSC_NULL, "Setting up PETSc KSP object.\n"); CHKERRB(mErr);
	KSP ksp;
	mErr = SNESGetKSP(mPetscSnes, &ksp); CHKERRB(mErr);
	mErr = KSPSetType(ksp, KSPGMRES); CHKERRB(mErr);	// preonly?
	mErr = KSPSetPCSide(ksp, PC_RIGHT); CHKERRB(mErr);	// for ForcingTermsEisenstatWalker
//	mErr = KSPSetFromOptions(ksp); CHKERRB(mErr);
	mErr = KSPSetConvergenceTest(ksp, ConvergenceTestWrapper, this, NULL); CHKERRB(mErr);

	// Set up PC
	mErr = PetscInfo(PETSC_NULL, "Setting up PETSc PC object.\n"); CHKERRB(mErr);
	PC pc;
//	mErr = PCCreate(PETSC_COMM_WORLD, &pc); CHKERRB(mErr);
//	mErr = KSPSetPC(ksp, pc); CHKERRB(mErr);
	mErr = KSPGetPC(ksp, &pc); CHKERRB(mErr);
	mErr = PCSetType(pc, PCLU); CHKERRB(mErr);
	mErr = PCFactorSetMatOrderingType(pc, MATORDERINGNATURAL); CHKERRB(mErr);
//	mErr = PCSetFromOptions(pc); CHKERRB(mErr);

	// SNESSetFromOptions sets SNES, KSP, and PC objects from command line options.
	// Call after the KSP and PC blocks above, to override defaults set there.
	mErr = SNESSetFromOptions(mPetscSnes); CHKERRB(mErr);

	// Get KSP tolerances after SNESSetFromOptions()
	mErr = KSPGetTolerances(ksp, &mKspInitialRelativeTolerance, &mKspInitialAbsoluteTolerance, &mKspInitialDivergenceTolerance, &mKspInitialMaxIterationCount); CHKERRB(mErr);

	// Set up preconditioner
	delete mpOwnedPreco;
	mpOwnedPreco = CreatePreconditionerFromOptions();
	if (mpOwnedPreco == NULL) return false;
	mErr = PetscInfo1(PETSC_NULL, "Setting '%s' preconditioner matrix.\n", mpOwnedPreco->GetName().c_str()); CHKERRB(mErr);
	mpPreco = mpOwnedPreco;

	// Set up forcing terms
	delete mpOwnedForcing;
	mpOwnedForcing = CreateForcingTermsFromOptions();
	if (mpOwnedForcing == NULL) return false;
	mErr = PetscInfo1(PETSC_NULL, "Setting '%s' forcing terms.\n", mpOwnedForcing->GetName().c_str()); CHKERRB(mErr);
	mpForcing = mpOwnedForcing;

	mConvergedReason = SNES_CONVERGED_ITERATING;
	mFunctionNorm = 0;
	mNonLinearIterationCount = 0;
	mLinearIterationCount = 0;
	mFunctionEvaluationCount = 0;

	mIsInitialized = true;
	return true;
} // NonLinearSolver::Initialize


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::SetPreconditioner"
bool NonLinearSolver::SetPreconditioner(Preconditioner& rPreco)
{
	if (!mIsInitialized) return false;
	mErr = PetscInfo1(PETSC_NULL, "Setting '%s' preconditioner matrix.\n", rPreco.GetName().c_str()); CHKERRB(mErr);

	// Set the preconditioner
	delete mpOwnedPreco;
	mpOwnedPreco = NULL;
	mpPreco = &rPreco;

	// Reset PC object so KSP knows it is a new preconditioner
	KSP ksp;
	mErr = SNESGetKSP(mPetscSnes, &ksp); CHKERRB(mErr);
	PC pc;
	mErr = KSPGetPC(ksp, &pc); CHKERRB(mErr);
	mErr = PCReset(pc); CHKERRB(mErr);

	return true;
} // NonLinearSolver::SetPreconditioner


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::SetForcingTerms"
bool NonLinearSolver::SetForcingTerms(ForcingTerms& rForcing)
{
	if (!mIsInitialized) return false;
	mErr = PetscInfo1(PETSC_NULL, "Setting '%s' forcing terms.\n", rForcing.GetName().c_str()); CHKERRB(mErr);

	// Set the forcing term strategy
	delete mpOwnedForcing;
	mpOwnedForcing = NULL;
	mpForcing = &rForcing;

	return true;
} // NonLinearSolver::SetForcingTerms


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::Solve"
bool NonLinearSolver::Solve(NonLinearProblem& rProblem)
{
	if (!mIsInitialized) return false;

	// Set non-linear problem
	mpProblem = &rProblem;

	// Get relevant vectors and matrices
	Vec pX = mpProblem->GetSolutionVector();
	Vec pF = mpProblem->GetFunctionVector();
	Mat pJ = mpProblem->GetJacobianMatrix();

	// Initialize the preconditioner
	// :WARNING: initialize preconditioner after calling mpProblem->GetJacobianMatrix(),
	//           in case that call creates matrix data needed for the preconditioner.
	mpPreco->Initialize(rProblem);

	Mat pM = mpPreco->GetMatrix();

	// Set up non-linear solver
	mErr = SNESSetFunction(mPetscSnes, pF, CalculateFunctionWrapper, this); CHKERRB(mErr);
	mErr = SNESSetJacobian(mPetscSnes, pJ, pM, CalculateJacobianWrapper, this); CHKERRB(mErr);

	// Get nonlinear residual tolerance
	PetscReal atol, rtol;
	mErr = SNESGetTolerances(mPetscSnes, &atol, &rtol, PETSC_NULL, PETSC_NULL, PETSC_NULL); CHKERRB(mErr);
	mErr = VecNorm(pF, NORM_2, &mFunctionNorm); CHKERRB(mErr);
	if (rtol > 0) rtol *= mFunctionNorm;	// rtol = absolute tolerance based on the relative tolerance
	if (rtol > atol) atol = rtol;			// atol = the significant absolute tolerance

	// Set forcing terms
	mpForcing->Initialize(mKspInitialRelativeTolerance, mKspInitialAbsoluteTolerance, mKspInitialDivergenceTolerance, mKspInitialMaxIterationCount, atol);

	// Solve the non-linear system of equations
	mErr = PetscInfo(PETSC_NULL, "Calling SNESSolve.\n"); CHKERRB(mErr);
	mErr = SNESSolve(mPetscSnes, PETSC_NULL, pX); CHKERRB(mErr);
	mErr = PetscInfo(PETSC_NULL, "SNESSolve done.\n"); CHKERRB(mErr);

	// Set some solver variables
	mErr = SNESGetConvergedReason(mPetscSnes, &mConvergedReason);
	mErr = VecNorm(pF, NORM_2, &mFunctionNorm); CHKERRB(mErr);
	mErr = SNESGetIterationNumber(mPetscSnes, &mNonLinearIterationCount); CHKERRB(mErr);
	mErr = SNESGetLinearSolveIterations(mPetscSnes, &mLinearIterationCount); CHKERRB(mErr);
	mErr = SNESGetNumberFunctionEvals(mPetscSnes, &mFunctionEvaluationCount); CHKERRB(mErr);

	// Print some data
//	mErr = PetscPrintf(PETSC_COMM_WORLD, "Convergence reason: %d\n", mConvergedReason); CHKERRB(mErr);
//	mErr = PetscPrintf(PETSC_COMM_WORLD, "Norm of F is %A after %d/%d iterations\n", GetFunctionNorm(), GetNonlinearIterationCount(), GetLinearIterationCount()); CHKERRB(mErr);

//	mErr = PCView(pc, PETSC_VIEWER_STDOUT_WORLD); CHKERRB(mErr);

	return true;
} // NonLinearSolver::Solve


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::CreatePreconditionerFromOptions"
/*static*/ Preconditioner* NonLinearSolver::CreatePreconditionerFromOptions(void)
{
	// Get the preconditioner matrix type option string
	std::string option;
	PetscErrorCode err = PetscUtils::GetOptionsString(msPetscOptionPrecoMatrixType, option, msPetscOptionMaxLength);
	if (err != 0) return NULL;

	// Set up the preconditioner object
	Preconditioner* pPreco = NULL;
	if (option == "")	// :NOTE: option string is empty if there was no -pc_mat_type option specified.
	{
		pPreco = new PreconditionerJacobian;
	}
	else if (option == "jacobian")
	{
		pPreco = new PreconditionerJacobian;
	}
	else if (option == "initialjacobian")
	{
		pPreco = new PreconditionerInitialJacobian;
	}
	else
	{
		PetscError(PETSC_COMM_SELF,__LINE__,PETSC_FUNCTION_NAME,__FILE__,__SDIR__,PETSC_ERR_ARG_UNKNOWN_TYPE,PETSC_ERROR_INITIAL,"Unable to find requested preconditioner matrix type %s",option.c_str());
		return NULL;
	}
	return pPreco;
} // NonLinearSolver::CreatePreconditionerFromOptions


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::CreateForcingTermsFromOptions"
/*static*/ ForcingTerms* NonLinearSolver::CreateForcingTermsFromOptions(void)
{
	// Get the forcing terms type option string
	std::string option;
	PetscErrorCode err = PetscUtils::GetOptionsString(msPetscOptionForcingTermsType, option, msPetscOptionMaxLength);
	if (err != 0) return NULL;

	// Set up the forcing terms object
	ForcingTerms* pForc = NULL;
	if (option == "")
	{
		pForc = new ForcingTermsFixed;
	}
	else if (option == "fixed")
	{
		pForc = new ForcingTermsFixed;
	}
	else if (option == "dembo")
	{
		pForc = new ForcingTermsDemboSteihaug;
	}
	else if (option == "eisenstat")
	{
		pForc = new ForcingTermsEisenstatWalker;
	}
	else if (option == "iterations")
	{
		pForc = new ForcingTermsIterations;
	}
	else
	{
		PetscError(PETSC_COMM_SELF,__LINE__,PETSC_FUNCTION_NAME,__FILE__,__SDIR__,PETSC_ERR_ARG_UNKNOWN_TYPE,PETSC_ERROR_REPEAT,"Unable to find requested forcing terms type %s",option.c_str());
		return NULL;
	}
	return pForc;
} // NonLinearSolver::CreateForcingTermsFromOptions


//============================= ACCESS     ===================================


bool NonLinearSolver::IsConverged(void) { return (mConvergedReason > 0); }
double NonLinearSolver::GetFunctionNorm(void) { return mFunctionNorm; }
int NonLinearSolver::GetNonlinearIterationCount(void) { return mNonLinearIterationCount; }
int NonLinearSolver::GetLinearIterationCount(void) { return mLinearIterationCount; }
int NonLinearSolver::GetFunctionEvaluationCount(void) { return mFunctionEvaluationCount; }


/////////////////////////////// PRIVATE    ///////////////////////////////////


//============================= WRAPPERS   ===================================


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::CalculateFunctionWrapper"
/*static*/ PetscErrorCode NonLinearSolver::CalculateFunctionWrapper(SNES snes, Vec pX, Vec pF, void* pCtx)
{
	PetscErrorCode err = PetscInfo(PETSC_NULL, "Function calculation wrapper called.\n"); CHKERRQ(err);
	NonLinearSolver* p = static_cast<NonLinearSolver*>(pCtx);
	if (p == NULL || p->mpProblem == NULL) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GENERAL, "Invalid user context");
	err = p->mpProblem->CalculateFunctionVector(snes, pX, pF); CHKERRQ(err);
	return 0;
} // NonLinearSolver::CalculateFunctionWrapper


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::CalculateJacobianWrapper"
/*static*/ PetscErrorCode NonLinearSolver::CalculateJacobianWrapper(SNES snes, Vec pX, Mat* ppJ, Mat* ppM, MatStructure* pFlag, void* pCtx)
{
	PetscErrorCode err = PetscInfo(PETSC_NULL, "Jacobian calculation wrapper called.\n"); CHKERRQ(err);
	NonLinearSolver* p = static_cast<NonLinearSolver*>(pCtx);
	if (p == NULL || p->mpProblem == NULL || p->mpPreco == NULL) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GENERAL, "Invalid user context");
	err = p->mpProblem->CalculateJacobianMatrix(snes, pX, ppJ); CHKERRQ(err);
	Mat pJ = *ppJ;
	err = p->mpPreco->Calculate(snes, pX, pJ, ppM, pFlag); CHKERRQ(err);
	return 0;
} // NonLinearSolver::CalculateJacobianWrapper


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::ConvergenceTestWrapper"
/* static */ PetscErrorCode NonLinearSolver::ConvergenceTestWrapper(KSP ksp, PetscInt it, PetscReal rnorm, KSPConvergedReason* pReason, void* pCtx)
{
	PetscErrorCode err = PetscInfo(PETSC_NULL, "Convergence test wrapper called.\n"); CHKERRQ(err);
	NonLinearSolver* p = static_cast<NonLinearSolver*>(pCtx);
	if (p == NULL || p->mpForcing == NULL) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GENERAL, "Invalid user context");
	err = p->mpForcing->TestConvergence(ksp, it, rnorm, pReason); CHKERRQ(err);
	return 0;
} // NonLinearSolver::ConvergenceTestWrapper


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::PCSetUpWrapper"
/* static */ PetscErrorCode NonLinearSolver::PCSetUpWrapper(PC pc)
{
	PetscErrorCode err = PetscInfo(PETSC_NULL, "PCSetup wrapper called.\n"); CHKERRQ(err);
	void* pCtx;
	err = PCShellGetContext(pc, &pCtx); CHKERRQ(err);
	NonLinearSolver* p = static_cast<NonLinearSolver*>(pCtx);
	if (p == NULL || p->mpPrecoShell == NULL) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GENERAL, "Invalid user context");
	err = p->mpPrecoShell->SetUp(pc); CHKERRQ(err);
	return 0;
} // NonLinearSolver::PCSetUpWrapper


#undef __FUNCT__
#define __FUNCT__ "NonLinearSolver::PCApplyWrapper"
/* static */ PetscErrorCode NonLinearSolver::PCApplyWrapper(PC pc, Vec x, Vec y)
{
	PetscErrorCode err = PetscInfo(PETSC_NULL, "PCApply wrapper called.\n"); CHKERRQ(err);
	void* pCtx;
	err = PCShellGetContext(pc, &pCtx); CHKERRQ(err);
	NonLinearSolver* p = static_cast<NonLinearSolver*>(pCtx);
	if (p == NULL || p->mpPrecoShell == NULL) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GENERAL, "Invalid user context");
	err = p->mpPrecoShell->Apply(pc, x, y); CHKERRQ(err);
	return 0;
} // NonLinearSolver::PCApplyWrapper


NONLINEARSOLVER_NAMESPACE_END
