/*$Id: ntr.c 1.98 05/05/10 15:21:39-05:00 sarich@zorak.(none) $*/

/* Note that this implementation will soon be changed.  It currently
accesses data in qcg.h -- this is bad! */
#include "tao_solver.h"

#ifdef TAO_USE_PETSC

#include "ntr.h"                /*I "tao_solver.h" I*/
#include "petscksp.h"
#include "src/ksp/ksp/kspimpl.h"
#include "src/ksp/ksp/impls/qcg/qcg.h"
#include "src/petsctao/linearsolver/taolinearsolver_petsc.h"

static int TaoDestroy_NTR(TAO_SOLVER,void *);

/*
    TaoSolve_NTR - Implements Newton's Method with a trust region approach 
    for solving unconstrained minimization problems.  

    The basic algorithm is taken from MINPACK-2 (dstrn).

    TaoSolve_NTR computes a local minimizer of a twice differentiable function
    f  by applying a trust region variant of Newton's method.  At each stage 
    of the algorithm, we us the prconditioned conjugate gradient method to
    determine an approximate minimizer of the quadratic equation

	     q(s) = g^T * s + .5 * s^T * H * s

    subject to the Euclidean norm trust region constraint

	     || D * s || <= delta,

    where delta is the trust region radius and D is a scaling matrix.
    Here g is the gradient and H is the Hessian matrix.

    Note:  TaoSolve_NTR MUST use the iterative solver KSPQCG; thus, we
           set KSPQCG in this routine regardless of what the user may have
           previously specified.
*/
#undef __FUNCT__  
#define __FUNCT__ "TaoSolve_NTR"
static int TaoSolve_NTR(TAO_SOLVER tao,void *ntrptr)
{
  TAO_NTR            *neP = (TAO_NTR *)ntrptr;
  int                info;
  double             xnorm, gnorm, max_val, ftrial, f, gdx,delta, zero = 0.0;
  TaoTruth           newton;
  /* Vec              X, G,  S, Xtrial; */
  TaoVec             *xx, *gg, *ss=neP->S, *xtrial=neP->Xtrial;
  TaoMat             *HH;
  KSP                pksp;
  KSP_QCG            *qcgP;
  TaoTerminateReason reason=TAO_CONTINUE_ITERATING;
  KSPConvergedReason kreason;
  TaoTruth           success;
  int                iter=0;
  TaoLinearSolver *tls;

  TaoFunctionBegin;

  info=TaoGetSolution(tao,&xx);CHKERRQ(info);
  info=TaoGetGradient(tao,&gg);CHKERRQ(info);
  info=TaoGetHessian(tao,&HH);CHKERRQ(info);

  info = TaoGetInitialTrustRegionRadius(tao,&delta);CHKERRQ(info);  /* trust region radius */

  info = xx->Norm2(&xnorm);CHKERRQ(info);              /* xnorm = || xx || */
  info = TaoComputeMeritFunctionGradient(tao,xx,&f,gg);CHKERRQ(info);
  info = gg->Norm2(&gnorm);CHKERRQ(info);               /* gnorm = || gg || */

  info = TaoGetLinearSolver(tao,&tls); CHKERRQ(info); 
  pksp=((TaoLinearSolverPetsc *)tls)->GetKSP();

  info = KSPSetType(pksp,KSPQCG);CHKERRQ(info);
  info = PetscLogInfo((tao,"TaoSolve_NTR: setting KSPType = KSPQCG\n")); CHKERRQ(info);
  qcgP = (KSP_QCG *) pksp->data;

  while (1) {

    info = TaoMonitor(tao,iter++,f,gnorm,0.0,delta,&reason);CHKERRQ(info);
    if (reason!=TAO_CONTINUE_ITERATING) break;

    newton        = TAO_FALSE;
    neP->success  = 0;
    info = TaoComputeHessian(tao,xx,HH);CHKERRQ(info);

    if (iter == 1) {			/* Initialize delta */
      if (delta <= 0) {
        if (xnorm > zero) delta = neP->factor1*xnorm;
        else {
	  info = TaoGetInitialTrustRegionRadius(tao,&delta);CHKERRQ(info);
	}
        info = HH->Norm1(&max_val);
        if (info == PETSC_ERR_SUP) {
          info = PetscLogInfo((tao,"TaoSolve_NTR: Initial delta computed without matrix norm info\n")); CHKERRQ(info);
        } else {
          if (TaoAbsDouble(max_val)<1.e-14)SETERRQ(PETSC_ERR_PLIB,"Hessian norm is too small");
          delta = TaoMax(delta,gnorm/max_val);
        } 
        delta = TAO_INFINITY;
      } else { 
	info = TaoGetInitialTrustRegionRadius(tao,&delta);CHKERRQ(info);
      }
    }
    do {
      /* Minimize the quadratic to compute the step s */
      qcgP->delta = delta;

      info = TaoPreLinearSolve(tao,HH);CHKERRQ(info);
      //      info = TaoLinearSolve(tao,HH,gg,ss,&success);CHKERRQ(info);
      info = TaoMinQPTrust(tao, HH, gg, ss, delta, &success); CHKERRQ(info);

      info = KSPGetConvergedReason(pksp,&kreason);CHKERRQ(info);
      info = ss->Dot(gg,&gdx); CHKERRQ(info);
      //      kreason=KSP_DIVERGED_ITS;

      if (kreason != KSP_CONVERGED_QCG_NEG_CURVE && kreason != KSP_CONVERGED_QCG_CONSTRAINED) {
        newton = TAO_TRUE;  /* truncated Newton step */
      } else if (gdx<=0){
	if ((int)kreason < 0) SETERRQ(1,"Failure in SLESSolve");
      }
      info=PetscLogInfo((tao,"TaoSolve_NTR: %d: ltsnrm=%g, delta=%g, q=%g \n", 
		 iter-1, qcgP->ltsnrm, delta, qcgP->quadratic));CHKERRQ(info);

      info = xtrial->Waxpby(1.0,xx,-1.0,ss);CHKERRQ(info);

      info = xtrial->Norm2(&xnorm);CHKERRQ(info);
                           		               /* ftrial = f(Xtrial) */
      info = TaoComputeFunction(tao,xtrial,&ftrial);CHKERRQ(info);

      /* Compute the function reduction and the step size */
      neP->actred = f - ftrial;
      neP->prered = -qcgP->quadratic;

      /* Adjust delta for the first Newton step */
      if ((iter == 1) && (newton)) delta = TaoMin(delta,qcgP->ltsnrm);

      if (neP->actred < neP->eta1 * neP->prered) {  /* Unsuccessful step */

         info = PetscLogInfo((tao,"TaoSolve_NTR: Rejecting step\n")); CHKERRQ(info);

         /* If iterate is Newton step, reduce delta to current step length */
         if (newton) {
           delta = qcgP->ltsnrm;
           newton = TAO_FALSE;
         }
         delta /= 4.0; 

      } else {          /* Successful iteration; adjust trust radius */

        neP->success = 1;
        info = PetscLogInfo((tao,"TaoSolve_NTR: Accepting step\n")); CHKERRQ(info);
        if (newton) {
           delta = sqrt(qcgP->ltsnrm*delta);
           if (neP->actred < neP->eta2 * neP->prered) delta /= 2.0;
        } else if (neP->actred < neP->eta2 * neP->prered)
           delta /= delta;
        else if ((neP->actred >= neP->eta3 * neP->prered) && 
           (neP->actred < neP->eta4 * neP->prered))
           delta *= 2.0;
        else if (neP->actred >= neP->eta4 * neP->prered)
           delta *= 4.0;
        else neP->sflag = 1;
      }

      neP->delta = delta;
      neP->xnorm=xnorm;

    } while (!neP->success);

    /* Question:  If (!neP->success && break), then last step was rejected, 
       but convergence was detected.  Should this update really happen? */
    info = xx->CopyFrom(xtrial);CHKERRQ(info);

    /* Note:  At last iteration, the gradient evaluation is unnecessary */
    /* Actually, already have the function, so just comput fjunk for now */
    info = TaoComputeFunctionGradient(tao,xx,&f,gg);CHKERRQ(info);

    info = gg->Norm2(&gnorm);CHKERRQ(info);
    /* need to set step variable */

  }

  TaoFunctionReturn(0);
}
/*------------------------------------------------------------*/
#undef __FUNCT__  
#define __FUNCT__ "TaoSetUp_NTR"
static int TaoSetUp_NTR(TAO_SOLVER tao, void *solver)
{
  int     n,info;
  TaoVec  *xx;
  TaoMat *HH;
  TaoLinearSolver *ksp;
  TAO_NTR *ctx = (TAO_NTR *)solver;

  TaoFunctionBegin;
  info = TaoGetSolution(tao,&xx);CHKERRQ(info);

  info=xx->Clone(&ctx->S); CHKERRQ(info);
  info=xx->Clone(&ctx->Xtrial); CHKERRQ(info);
  info=xx->Clone(&ctx->G); CHKERRQ(info);

  info = TaoSetLagrangianGradientVector(tao,ctx->G);CHKERRQ(info);
  info = TaoSetStepDirectionVector(tao,ctx->S);CHKERRQ(info);

  info = TaoGetHessian(tao,&HH); CHKERRQ(info);
  info = TaoCreateLinearSolver(tao,HH,220,0); CHKERRQ(info); 
  info = TaoGetLinearSolver(tao,&ksp); CHKERRQ(info); 
  info = xx->GetDimension(&n);CHKERRQ(info);
  info = ksp->SetTolerances(1e-8,1e-30,1e30,n);CHKERRQ(info);

  info = TaoCheckFGH(tao);CHKERRQ(info);

  TaoFunctionReturn(0);
}
/*------------------------------------------------------------*/
#undef __FUNCT__  
#define __FUNCT__ "TaoDestroy_NTR"
static int TaoDestroy_NTR(TAO_SOLVER tao,void *solver)
{
  TAO_NTR *ctx = (TAO_NTR *)solver;
  int  info;

  TaoFunctionBegin;

  info=TaoVecDestroy(ctx->G);CHKERRQ(info);
  info=TaoVecDestroy(ctx->S);CHKERRQ(info);
  info=TaoVecDestroy(ctx->Xtrial);CHKERRQ(info);
  info = TaoDestroyLinearSolver(tao);CHKERRQ(info);

  info = TaoSetLagrangianGradientVector(tao,0);CHKERRQ(info);
  info = TaoSetStepDirectionVector(tao,0);CHKERRQ(info);

  TaoFunctionReturn(0);
}
/*------------------------------------------------------------*/
#undef __FUNCT__  
#define __FUNCT__ "TaoSetOptions_NTR"
static int TaoSetOptions_NTR(TAO_SOLVER tao, void*solver)
{
  TAO_NTR    *ctx = (TAO_NTR *)solver;
  int        info;

  TaoFunctionBegin;

  info = TaoOptionsHead("Newton trust region method for unconstrained optimization");CHKERRQ(info);
  info = TaoOptionDouble("-tao_ntr_eta1","step is unsuccessful if actual reduction < eta1 * predicted reduction","",ctx->eta1,&ctx->eta1,0);CHKERRQ(info);

  info = TaoOptionDouble("-tao_ntr_eta2","","",ctx->eta2,&ctx->eta2,0);CHKERRQ(info);
  info = TaoOptionDouble("-tao_ntr_eta3","","",ctx->eta3,&ctx->eta3,0);CHKERRQ(info);
  info = TaoOptionDouble("-tao_ntr_eta4","","",ctx->eta3,&ctx->eta3,0);CHKERRQ(info);
  info = TaoOptionDouble("-tao_ntr_factor1","","",ctx->factor1,&ctx->factor1,0);CHKERRQ(info);  

  /* set the special KSP monitor for matrix-free application */
  /*
  info = TaoSetLinearSolverOptions(tao); CHKERRQ(info);
  info = OptionsHasName(tao->prefix,"-tao_mf_ksp_monitor",&flg); CHKERRQ(info);
  if (flg) {
  SLES sles;
    KSP ksp;
    info = SLESGetKSP(sles,&ksp);CHKERRQ(info);
    info = KSPSetMonitor(ksp,MatSNESMFKSPMonitor,TAO_NULL,0); CHKERRQ(info);
  }
  */

  info = TaoOptionsTail();CHKERRQ(info);
  TaoFunctionReturn(0);
}
/*------------------------------------------------------------*/

#undef __FUNCT__  
#define __FUNCT__ "TaoView_NTR"
static int TaoView_NTR(TAO_SOLVER tao,void*solver)
{
  TAO_NTR         *tr = (TAO_NTR *)solver;
  int             info;

  TaoFunctionBegin;

  info = TaoPrintStatement(tao,"  Linear solver minimizes quadratic over Trust Region: \n");CHKERRQ(info);

  info = TaoPrintDouble(tao,"  eta1=%g,",tr->eta1);CHKERRQ(info);
  info = TaoPrintDouble(tao," eta1=%g,",tr->eta2);CHKERRQ(info);
  info = TaoPrintDouble(tao," eta3=%g,",tr->eta3);CHKERRQ(info);
  info = TaoPrintDouble(tao," eta4=%g\n",tr->eta4);CHKERRQ(info);
  info = TaoPrintDouble(tao,"  factor1=%g\n",tr->factor1);CHKERRQ(info);

  TaoFunctionReturn(0);
}
/*------------------------------------------------------------*/
EXTERN_C_BEGIN
#undef __FUNCT__  
#define __FUNCT__ "TaoCreate_NTR"
int TaoCreate_NTR(TAO_SOLVER tao)
{
  TAO_NTR         *neP;
  int             info;

  TaoFunctionBegin;

  info = TaoNew(TAO_NTR,&neP); CHKERRQ(info);
  info = PetscLogObjectMemory(tao,sizeof(TAO_NTR)); CHKERRQ(info);

  info=TaoSetTaoSolveRoutine(tao,TaoSolve_NTR,(void*)neP); CHKERRQ(info);
  info=TaoSetTaoSetUpDownRoutines(tao,TaoSetUp_NTR,TaoDestroy_NTR); CHKERRQ(info);
  info=TaoSetTaoOptionsRoutine(tao,TaoSetOptions_NTR); CHKERRQ(info);
  info=TaoSetTaoViewRoutine(tao,TaoView_NTR); CHKERRQ(info);

  info = TaoSetMaximumIterates(tao,50); CHKERRQ(info);
  info = TaoSetTolerances(tao,1e-10,1e-10,0,0); CHKERRQ(info);
  info = TaoSetTrustRegionRadius(tao,1);CHKERRQ(info);

  //  tao->converged	= TaoConverged_Default;

  info = TaoSetTrustRegionTolerance(tao,1.0e-12);CHKERRQ(info);

  neP->delta 		= 0.0;
  neP->eta1		= 1.0e-4;
  neP->eta2		= 0.25;
  neP->eta3		= 0.50;
  neP->eta4		= 0.90;
  neP->factor1		= 1.0e-8;
  neP->actred		= 0.0;
  neP->prered		= 0.0;
  neP->success		= 0;
  neP->sflag		= 0;
  neP->xnorm            = 1.0e+20;
  /* Set default preconditioner to be Jacobi, to override KSP default. */
  /* This implementation currently requires a symmetric preconditioner. */

  /* info = KSPGetPC(tao->ksp,&pc);CHKERRQ(info);
  info = PCSetType(pc,PCJACOBI);CHKERRQ(info); */

  TaoFunctionReturn(0);
}
EXTERN_C_END


#endif

