/*--------------------------------------------------------------------------
 * solvenle.ox - code for solving nonlinear equations
 *
 *       (C) Jurgen Doornik 2002-
 *
 * It is not permitted to translate this source code into any other
 * computing language. This code may only be used as part of the
 * Ox system.
 *--------------------------------------------------------------------------*/
#include <oxstd.h>
#include <oxfloat.h>
#import  <maximize>

/*----------------------- line searches for SolveNLE -----------------------*/
// func		 in: function to do linesearch over
// avX		 in: address of n x 1 vector, current point
//			 out:n x 1 vector with new point
// avF   	 in: address of 1 x 1 or n x 1 vector, function value(s) at vX
//			 out:new function values at avX (unchanged if line search failed)
// dSlope	 in: double, g'Jp: slope of direction
// vDelta	 in: n x 1 vector, direction
// dEps1 	 in: double, tolerance for cubic part
// dEps2 	 in: double, tolerance for linear part
// adStep 	 in: address of double: initial steplength
//			 out: double, actual step length
//
// Returns
// integer:
// 		0: failed to improve
// 		1: normal convergence in line search
//
// Cubic part is similar to Numerical recipes algorithm lnsrch (Section 9.7, pp384-389)
static lineSearchNleCubic(const func, const avX, const avF, const dSlope,
	const vDelta, const dEps1, const dEps2, const adStep)
{
	decl ax1, a, b, disc, r1, r2, dstep, dstepprev, dstepmin, dstepabsmin,
		newstep, scale, fquad;
	decl x0 = avX[0], fi = avF[0], f0, f, fprev, ret_val, i;

	f0 = f = double(sumsqrc(vec(fi)) / 2);
	dstep = adStep[0]; 	// try full (Newton) step first.
	adStep[0] = 0;		// set to zero for failure

	a = norm(x0);
	b = norm(vDelta);

	if (b == 0 || dSlope == 0 || ismissing(vDelta))
		return 0;		// step direction is zero
	dstepmin = a == 0 ? dEps1 : dEps1 * a / b;
	dstepabsmin = a == 0 ? dEps2 : dEps2 * a / b;

	if (dSlope < 0)
	{
		for (i = 0, fquad = TRUE; dstep > dstepmin; ++i)
		{
			fprev = f;
			// evaluate function at this x0
			ret_val = func(&fi, x0 + dstep * vDelta);
			f = double(sumsqrc(vec(fi)) / 2);
	
			// in case of failure or lack of improvement: use linear search
			if (!ret_val || ismissing(f))
				break;

			if (f <= f0 + dEps1 * dstep * dSlope)
			{						// yes: Armijo condition satisfied
				adStep[0] = dstep;
				avX[0] = x0 + dstep * vDelta;
				avF[0] = fi;
				return 1; 			// success: sufficient function decrease.
			}
			if (f - fprev > 100 * (1 + fabs(fprev)))
			{						// this was a very bad step
				dstepprev = dstep;
				dstep /= 10;
				fquad = TRUE;
				continue;
			}
		 	// now enter backtracking stage
			if (fquad)	 			// yes: first backtrack uses quadratic
			{
				newstep = -dSlope / ( 2.0 * (f - f0 - dSlope) );
				fquad = FALSE;
			}
			else					// subsequent backtracks use cubic
			{	 
				r1 = (f     - f0 - dstep     * dSlope) / sqr(dstep);
				r2 = (fprev - f0 - dstepprev * dSlope) / sqr(dstepprev);
				a = (r1 - r2) / (dstep - dstepprev);
				b = (-dstepprev * r1 + dstep * r2) / (dstep - dstepprev);
				if (a == 0.0)
					newstep = -dSlope / (2.0 * b);
				else
				{
					disc = sqr(b) - 3.0 * a * dSlope;
					if (disc < 0.0)
						newstep = 0.5 * dstep;
					else if (b <= 0.0)
						newstep = (-b + sqrt(disc)) / (3.0 * a);
					else
						newstep = -dSlope / (b + sqrt(disc));
				}
			}
			if (newstep > 0.5 * dstep)
				newstep = 0.5 * dstep;
			else if (newstep < 0.1 * dstep)
				newstep = 0.1 * dstep;
			dstepprev = dstep;
			dstep = newstep;
		}
	}

	scale = 10.0;
	if (dSlope >= 0.0)			// yes: non-negative dSlope
		scale = -2.0;		 		// alternating (smaller) steps
	for (; ; dstep /= scale)
	{
		// evaluate function at this x0
		ret_val = func(&fi, x0 + dstep * vDelta);
		fprev = f;
		f = double(sumsqrc(vec(fi)) / 2);

		if (fabs(dstep) <= dstepabsmin)
		{
			adStep[0] = dstep;
			return 0;			// here: no success in linesearch
		}
		if (!ret_val || ismissing(f))	
			continue;			// yes: continue until successful evaluation

		if (f - f0 <= -dEps1 * fabs(dstep * dSlope))
		{						// yes: Armijo condition satisfied
			adStep[0] = dstep;
			avX[0] = x0 + dstep * vDelta;
			avF[0] = fi;
			return 1; 			// success: sufficient function decrease.
		}
	}
}

static lineSearchNle(const func, const avX, const avF, dSlope, vDelta,
	const dEps, const adStep, dSlopeAlt, vDeltaAlt)
{
	decl fstepok;
	adStep[0] = 1;			// initial steplength

	fstepok = lineSearchNleCubic(func, avX, avF, dSlope, vDelta, dEps,
			DBL_EPSILON, adStep);
			
	if (!fstepok && sizer(vDeltaAlt))
	{
		adStep[0] = 1;		// initial steplength
		if (dSlopeAlt < 0)	// desperate: try positive and negative directions
			dSlopeAlt = -dSlopeAlt;
		fstepok = lineSearchNleCubic(func, avX, avF, dSlopeAlt, vDeltaAlt,
			dEps, DBL_EPSILON, adStep);
	}
	return fstepok;
}
/*--------------------- END line searches for SolveNLE ---------------------*/

/*------------------------------- SolveGMRes -------------------------------*/
SolveGMRes(const funcAx, const vB, const vX, eps, m, ...)
{
	decl i, j, normvj, norm0, vj, h, hj, mv, n = rows(vB), y, g;
	decl args = va_arglist(), fusefuncjx = FALSE, funcjx;
	if (sizeof(args))
		funcjx = args[0], fusefuncjx = TRUE;
	
	if (eps == 0) eps = 1e-8;		// then use default eps
	norm0 = norm(vB);
	if (norm0 == 0)					// right-hand side is zero 
		return {zeros(n,1), 0};
	eps *= norm0;
	
	vj = vB - funcAx(vX);			// current residual

	norm0 = normvj = norm(vj, 2);
	if (norm0 == .NaN)
		return {constant(.NaN, n, 1), 0};
	if (norm0 <= eps)
		return {vX, norm0};

	if (m > n) m = n;
	h = zeros(m, m + 1);			// will hold upper Hessenberg matrix H'
	mv = <>;						// will hold orthogonal basis V[n][m]
	for (j = 0; j < m && normvj > eps; ++j)
	{
		vj /= normvj;				// normalize residual
		mv ~= vj;
		vj = fusefuncjx ? funcjx(vj, vX) : funcAx(vj);
		hj = vj'mv;					// hj is a row vector
		vj -= mv * hj';
		normvj = norm(vj, 2);
		h[j][] = hj ~ normvj;		// store in H'
	}
	g = norm0 ~ zeros(1, j);
	if (j < m) h = h[ : j - 1][ : j];
	olsr(g, h, &y);

	return {vX + mv * y', norm(g - y * h, 2)};
}
/*------------------------------- SolveGMRes -------------------------------*/

/*--------------------- tensor functions for SolveNLE ----------------------*/
static decl s_mTensorA_tf, s_mS_tf, s_vF_tf, s_vX_tf, s_func_tf;

static getTensor(const mFhist, const vF, const mS, const mJacS)
{
	if (columns(mFhist) <= 0)
		return <>;
	// find tensor matrix A, Schnabel & Frank (1984) Theorem 3.2
	decl scale = sumsqrc(mS), ml, vd;
	scale = scale .< 1e-20 .? 1 .: 1 ./ scale;
	decl mz = 2 * (mFhist - vF - mJacS);
	decl ms = mS .* sqrt(scale);
	decl mm = sqr(ms'ms);

	return (mz .* scale) * (invertgen(mm, 3) .* scale);
}
static tensorCorrection(const mFhist, const vF, const mXhist, const vX,
	const mJac, vP)
{
	decl ms = mXhist - vX, tensora, p, n = sizer(vF);

	tensora = getTensor(mFhist, vF, ms, mJac * ms);
	if (columns(tensora) == 0)
		return vP;

	// solve tensor problem: f + Jp + A.sqr(S'p)/2 = g, min g'g/2
	decl i, g, gjac, f, d, f0, fold, pold;
	// take some Gauss-Newton steps

	for (i = 0, pold = vP; i < 10; ++i)
	{
		g = vF + mJac * vP + (tensora * sqr(ms'vP)) / 2;
		f = double(sumsqrc(g) / 2);
		if (i == 0)
			f0 = f;
		else
		{
			if (fold <= f)				// yes: step in wrong direction
			{
				vP = pold;  f = fold;
				break;
			}
			if (f < 1e-10 || fold - f < 1e-2 * fabs(f0))
				break;					// yes: `weak' convergence
		}
		// derivative of tensor problem is: H = J + A(S' .* (S'p))
		// (derivative of g'g/2 is H'g; unfortunately large scale
		//  methods can easily compute Hg, but not H'g)
		gjac = mJac + tensora * (ms' .* (ms'vP));

		olsc(g, gjac, &d);
		fold = f;  pold = vP;
		vP -= d;
	}
	return vP;
}
static FuncAx_solveNLE(const vX)
{
	decl jacx;
	NumJacobianMul(s_func_tf, s_vX_tf, s_vF_tf, vX, TRUE, &jacx);
	if (columns(s_mTensorA_tf))
		jacx += (s_mTensorA_tf * sqr(s_mS_tf'vX)) / 2;
	return jacx;
}
static FuncJx_solveNLE(const vX, const vX0)
{
	decl jacx;
	NumJacobianMul(s_func_tf, s_vX_tf, s_vF_tf, vX, TRUE, &jacx);
	if (columns(s_mTensorA_tf))	// parentheses very important to avoid n x n matrix
		jacx += s_mTensorA_tf * ((s_mS_tf' .* (s_mS_tf'vX0))*vX);
		
	return jacx;
}
static tensorCorrectionCG(const mFhist, const vF, const mXhist, const vX,
	const func, vP, const dEps, const mxIter)
{
	decl n = sizer(vF);
	decl i, j, g, f, f0, fold, xold, df, dfprev, p, beta, dsteplen,
		fstepok, mjs, err;

	s_mS_tf = mXhist - vX;
	s_mTensorA_tf = <>;

	// evaluate JS by forwarde differences (remainder uses central diffs
	if (NumJacobianMul(func, vX, vF, s_mS_tf, FALSE, &mjs))
		s_mTensorA_tf = getTensor(mFhist, vF, s_mS_tf, mjs);
	if (columns(s_mTensorA_tf) == 0)
		return vP;
	mjs = 0;
	
	s_vF_tf = vF;
	s_func_tf = func;
	s_vX_tf = vX;

	// very ad hoc: iterating twice using nonlinear gmres
	[vP, err] = SolveGMRes(FuncAx_solveNLE, -vF, vP, dEps, columns(s_mTensorA_tf)    , FuncJx_solveNLE);
	[vP, err] = SolveGMRes(FuncAx_solveNLE, -vF, vP, dEps, 2 * columns(s_mTensorA_tf), FuncJx_solveNLE);

	s_mS_tf = s_mTensorA_tf = s_vF_tf = s_func_tf = s_vX_tf = 0;
	return vP;
}
/*------------------- END tensor functions for SolveNLE --------------------*/

/*------------------------------ SolveNLE ----------------------------------*/
static step_solveNLE(const func, const mFhist, const vF, const mXhist, const vX,
	const mJac)
{
	decl mh, mr, mp, pp, pn, alpha, beta, scale, pc, gg;
	
	// solve Jp=-f, allowing for singularity
	decqr(mJac, &mh, &mr, &mp);				
	pp = decqrmul(mh, vF);  				// JP=QR, pp=Q'f
	pn = solvelu(0, mr, 0, -pp)[mp[0][]];	// Newton step

	// Compute cauchy point -||J'f||/||JJ'f||*J'f where ||a||=sumsqrc(a)
	gg = mr'pp;								// = P'J'f
	scale = norm(gg);
	alpha = double(sumsqrc(gg / scale));
	beta = double(sumsqrc(mr * (gg / scale)));
	gg = gg[mp[0][]];						// nb: beta needs not resorted gg
	pc = -alpha / beta * gg;				// cauchy step

	if (columns(mFhist))
	{
		pn = tensorCorrection(mFhist, vF, mXhist, vX, mJac, pn);
	}
	// gradient' = f'J, slope = f'Jp_n, steepest descent slope f'Jp_c
	return {pn, pc, double(gg'pn), double(gg'pc)};
}
static	decl s_time = 0;


static step_SolveNLE_Large(const func, const mFhist, const vF,
	const mXhist, const vX, const dEps1, const dEps2, const mxIter)
{
	decl pn, pc, slopen, slopes, err;

	pn = zeros(sizer(vF), 1);

	s_vF_tf = vF;  s_func_tf = func;  s_vX_tf = vX;
	// linear gmres, twice (which is quite ad hoc)
	[pn, err] = SolveGMRes(FuncAx_solveNLE, -vF, pn, dEps1, 10);
	[pn, err] = SolveGMRes(FuncAx_solveNLE, -vF, pn, dEps1, mxIter);
	s_vF_tf = s_func_tf = s_vX_tf = 0;
	
	pc = pn;					// use gmres-Newton as alternative
		
	// iterates on linear system if tensor available
	pn = tensorCorrectionCG(mFhist, vF, mXhist, vX, func, pn, dEps2, mxIter);

	NumJacobianMul(func, vX, vF, pn, FALSE, &slopen);
	slopen = double(vF'slopen);
	NumJacobianMul(func, vX, vF, pc, FALSE, &slopes);
	slopes = double(vF'slopes);

	return {pn, pc, double(slopen), double(slopes)};
}

static getparams_solveNLE(const args)
{
	decl mode = -1, fnjac = 0, deps1, deps2, mxiter, iprint,
		mxit_inner = -1, cargs = sizeof(args);
	// first process default arguments settings
	[mxiter, iprint] = GetMaxControl();
	[deps1, deps2] = GetMaxControlEps();

	if (cargs > 0)  			   mode 	= 	args[0];
	if (cargs > 1)  			   fnjac 	= 	args[1];
	if (cargs > 2 && args[2] > 0)  deps1 	= 	args[2];
	if (cargs > 3 && args[3] > 0)  deps2 	= 	args[3];
	if (cargs > 4 && args[4] >= 0) mxiter 	= 	args[4];
	if (cargs > 5 && args[5] >= 0) iprint 	= 	args[5];
	if (cargs > 6 && args[6] >= 0) mxit_inner = args[6];
	return {mode, fnjac, deps1, deps2, mxiter, iprint, mxit_inner};
}
SolveNLE(const func, const avX, ...)
{
	// first process default arguments settings
	decl mode, funcjac, deps1, deps2, mxiter, iprint, bcompact, mxit_inner;
	[mode, funcjac, deps1, deps2, mxiter, iprint, mxit_inner]
		= getparams_solveNLE(va_arglist());
	
	// start of SolveNLE
	decl fstepok, fresetj, fweak = FALSE;
	decl itno, n, ret_val = MAX_NOCONV, resetit = -1, c, ctensor;
	decl fcrit, fold, test, dtol, fcritinit, dsteplen = 1, steptol, slope,
		slopec, normf;
	decl fvec, xold, fvecold, g, pn, pc;
	decl mjac, fvechist = <>, xhist = <>;
	
	if (sizec(avX[0]) > 1)
		avX[0] = vec(avX[0]);	// ensure row vector
	n = sizer(avX[0]);
	if (n <= 0)
		return MAX_NOCONV;

	if (mode == -1)
		mode = n < 80 ? 1 : 3;
	if (mode == 0 && !isfunction(funcjac))
	{
		println("SolveNLE error: need to supply function for Jacobian");
		return MAX_NOCONV;
	}
	if (mode < 0 || mode > 3)
	{
		println("SolveNLE error: invalid mode argument");
		return MAX_NOCONV;
	}
	ctensor = int(round(sqrt(n)));
	if (ctensor >= n) ctensor = n - 1;
	if (ctensor >= 5) ctensor = 5;
	if (mxit_inner < 0) mxit_inner = max(50, int(10 * log10(n)));

	dtol = deps1 * 1e-3;
	steptol = 100 * DBL_EPSILON;

	func(&fvec, avX[0]);		// evaluate function
	fcritinit = fcrit = double(sumsqrc(vec(fvec))) / 2;
	normf = norm(fvec);

	for (itno = 0, fresetj = TRUE; ; ++itno)
	{
		// get Jacobian; mode=0: analytical, 1: numerical, 2: Broyden, 3: large scale
		if (mode == 0)
		{	funcjac(&mjac, avX[0]);
		}
		else if ( (fresetj && mode == 2) || mode == 1)
		{								// numerical or Broyden/Column reset
			NumJacobianEx(func, avX[0], fvec, FALSE, &mjac);
			resetit = itno;
		}
		else if (mode == 2)			// use Broyden update
		{	g = (fvec - fvecold - mjac * pn);
			g = fabs(g) .<= DBL_EPSILON * (fabs(fvec) + fabs(fvecold)) .? 0 .: g;
			mjac += g * (pn' / (pn'pn));
		}
		fresetj = FALSE;

		xold = avX[0]; 
		fold = fcrit;
		fvecold = fvec;
		if (iprint > 0 && (imod(itno, iprint) == 0 || ret_val != MAX_NOCONV) )
		{
			print("it", "%-6d", itno, "f'f/2=", "%#13.7g", fold, " ||f||=", "%#11.5g", normf);
			if (itno)
			{	print(" slope=", "%#12.5g", slope * dsteplen);
				if (dsteplen != 1)
					print(" step=", "%.4g", dsteplen);
			}
			print("\n");
	        if (ret_val != MAX_NOCONV)
    	        println("SolveNLE(", mode, "): ", MaxConvergenceMsg(ret_val));
		}
		
		if (ret_val != MAX_NOCONV)
			break;

		if (mode == 3)				// large scale Newton
			[pn, pc, slope, slopec] = step_SolveNLE_Large(func, fvechist,
				fvec, xhist, xold, deps1, deps2, mxit_inner);
		else
			[pn, pc, slope, slopec] = step_solveNLE(func, fvechist, fvec, xhist, xold, mjac);

		if (ctensor && imod(itno,2) == 0)
		{							// only every 2nd for tensor (ad hoc decision)
			fvechist ~= fvec;
			xhist ~= xold;
			if (columns(xhist) > ctensor)
				xhist = xhist[][1:], fvechist = fvechist[][1:];
		}
		// if slope is not descent but slopec is: start with that instead
		if (slope >= -1e-14 && slopec < -1e-14)
			fstepok = lineSearchNle(func, avX, &fvec, slopec, pc, deps1,
				&dsteplen, slope, pn);
		else
			fstepok = lineSearchNle(func, avX, &fvec, slope, pn, deps1,
				&dsteplen, slopec, pc);

		pn = avX[0] - xold;			// actual stepsize
		fcrit = double(sumsqrc(vec(fvec))) / 2;
		normf = norm(fvec);
		fweak = normf < deps2;

		if (ismissing(fcrit))
			ret_val = MAX_FUNC_FAIL;
		else if (normf < dtol)		// check convergence on function values.
			ret_val = MAX_CONV;
		else if (itno >= mxiter)
			ret_val = MAX_MAXIT;
		else if (!fstepok || fcrit - fold == 0 || norm(pn) <= steptol * max(fabs(avX[0]), 1) )
		{							// failed in linesearch
			if (mode == 2 && resetit != itno)
				fresetj = TRUE, ret_val = MAX_NOCONV;
			else if (fweak)			// weak convergence on function values.
				ret_val = MAX_WEAK_CONV;
			else
				ret_val = MAX_LINE_FAIL;
		}
	}
	return ret_val;
}
/*---------------------------- END SolveNLE --------------------------------*/
