#include "BSprivate.h"

/*@ BSpar_isolve - Solve a symmetric indefinite system of equations
                   using symmlq preconditioned by one of several
                   preconditioners.

    Input Parameters:
.   A - a sparse matrix
.   fact_A - the incomplete factored version of A, if any (NULL if not exist)
.   comm_A - the communication structure for A
.   B - a sparse matrix 
.   comm_B - the communication structure for B
.   in_rhs - the rhs
.   shift - the shift to multiply B by
.   pre_option - the preconditioner to use
$     PRE_ICC: incomplete Cholesky factorization of A
$     PRE_ILU: incomplete LU factorization of A
$     PRE_SSOR: Successive over relaxation (using just A)
$     PRE_BJACOBI: Block Jacobi (using just A)
.   residual - the final computed residual
.   procinfo - the usual processor stuff

    Output Parameters:
.   out_x - the solution vector

    Returns:
    The number of iterations.

    Notes:
    The system solved is (A-shift*B)out_x = in_rhs.
    The preconditioners must be computed prior to calling BSpar_isolve.
    For more information on the preconditioners, see the manual.

$   The following are now specified in the context:
$   err_tol - the tolerance to which to solve the problem
$     stop if the estimated norm of the residual divided by
$     the norm of the rhs is less than err_tol
$   max_iter - the maximum number of iterations to take
$   guess - if TRUE, then initialize out_x to 0, otherwise
$     the program assumes that out_x contains an initial
$     guess

 @*/
int BSpar_isolve(BSpar_mat *A, BSpar_mat *fact_A, BScomm *comm_A,
	BSpar_mat *B, BScomm *comm_B, FLOAT *in_rhs, FLOAT *out_x,
	FLOAT shift, FLOAT *residual, BSprocinfo *procinfo)
{
	/* temporary space to hold an incoming vector */
	int	i;
	int	cur_step, cur_phase;
	int	done, conv;
	FLOAT	*resid, *z, *v, *w, *r2;
	FLOAT	cg_beta, cg_alpha;
	FLOAT	obeta;
	FLOAT	gbar, dbar;
	FLOAT	rhs1, rhs2;
	FLOAT	snprod;
	FLOAT	beta1;
	FLOAT	bstep;
	FLOAT	ynorm2;
	FLOAT	tnorm;
	FLOAT	bestnm;
	int	cgpt;
	FLOAT bnorm;
	FLOAT *x;
	FLOAT *rhs;
	FLOAT *temp1, *temp2;
	FLOAT	tval;
	int pre_option, max_iter, guess;
	FLOAT err_tol;

	if ((procinfo->print) && (PSISROOT(procinfo))) {
		BSctx_print(procinfo);
	}

	pre_option = procinfo->preconditioner;
	err_tol = procinfo->residual_tol;
	max_iter = procinfo->max_iterations;
	guess = procinfo->guess;

	/* reorder the rhs */
	MY_MALLOCN(rhs,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	BSperm_dvec(in_rhs,rhs,A->perm); CHKERRN(0);

	/* allocate space for x */
	MY_MALLOCN(x,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);

	/* allocate space for symmlq vectors */
	MY_MALLOCN(resid,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	MY_MALLOCN(v,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	MY_MALLOCN(r2,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	MY_MALLOCN(w,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	MY_MALLOCN(z,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	MY_MALLOCN(temp1,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);
	MY_MALLOCN(temp2,(FLOAT *),sizeof(FLOAT)*A->num_rows,1);

	/* form the initial guess */
	if (guess) {
		for (i=0;i<A->num_rows;i++) {
			x[i] = 0.0;
		}
	} else {
		BSperm_dvec(out_x,x,A->perm);
	}

	/* scale the rhs */
	for (i=0;i<A->num_rows;i++) {
		tval = sqrt(A->scale_diag[i]);
		rhs[i] /= tval;
		x[i] *= tval;
	}

	/* get the norm of B */
	bnorm = BSpar_ip(A->num_rows,rhs,rhs,procinfo); CHKERRN(0);
	bnorm = sqrt(bnorm);
	done = FALSE;
	cur_step = 0;
	cur_phase = 0;
	if (bnorm == 0.0) done = TRUE;
	conv = FALSE;
	while ((!done) && (cur_step < max_iter)) {
		switch(BSpar_symmlq(A->num_rows,rhs,x,resid,z,v,w,r2,&cg_beta,&cg_alpha,
			&obeta,&gbar,&dbar,&rhs1,&rhs2,
			&bstep,&snprod,&beta1,&tnorm,&ynorm2,&bestnm,&cgpt,conv,
			&cur_step,&cur_phase,procinfo)) {
			case SYMM_MATVECXR: {
				BStri_mult(A,comm_A,B,comm_B,x,resid,temp1,temp2,shift,1,
					procinfo); CHKERRN(0);
				break;
			}
			case SYMM_MATVECVZ: {
				/* check for convergence first */
				if ((cur_step > 0) && (bestnm/bnorm < err_tol)) conv = TRUE;
				if (!conv) {
					BStri_mult(A,comm_A,B,comm_B,v,z,temp1,temp2,shift,1,
						procinfo); CHKERRN(0);
				}
				break;
			}
			case SYMM_MSOLVE1: {
				for (i=0;i<A->num_rows;i++) {
					z[i] = resid[i];
				}
				BStri_solve(A,fact_A,comm_A,z,pre_option,1,procinfo); CHKERRN(0);
				break;
			}
			case SYMM_MSOLVE2: {
				for (i=0;i<A->num_rows;i++) {
					z[i] = r2[i];
				}
				BStri_solve(A,fact_A,comm_A,z,pre_option,1,procinfo); CHKERRN(0);
				break;
			}
			case SYMM_MSOLVEB: {
				for (i=0;i<A->num_rows;i++) {
					z[i] = rhs[i];
				}
				BStri_solve(A,fact_A,comm_A,z,pre_option,1,procinfo); CHKERRN(0);
				break;
			}
			case SYMM_DONE: {
				done = TRUE;
				break;
			}
			default: {
				return(-cur_step);
			}
		}
		CHKERRN(0);
	}

	BStri_mult(A,comm_A,B,comm_B,x,resid,temp1,temp2,shift,1,procinfo); 
	CHKERRN(0);
	for (i=0;i<A->num_rows;i++) {
		resid[i] = rhs[i] - resid[i];
	}
	(*residual) = BSpar_ip(A->num_rows,resid,resid,procinfo); CHKERRN(0);
	if (bnorm != 0.0) {
		(*residual) = sqrt((*residual))/bnorm;
	} else {
		(*residual) = sqrt((*residual));
	}

	MY_FREE(temp1);
	MY_FREE(temp2);
	MY_FREE(z);
	MY_FREE(v);
	MY_FREE(w);
	MY_FREE(r2);
	MY_FREE(resid);

	/* Rescale X */
	for (i=0;i<A->num_rows;i++) {
		x[i] /= sqrt(fabs(A->scale_diag[i]));
	}

	/* reorder the solution vector */
	BSiperm_dvec(x,out_x,A->perm); CHKERRN(0);
	MY_FREE(rhs);
	MY_FREE(x);

	/* return the number of iterations */
	return(cur_step);
}