#include "pargrid.h"

/*+ worker - Solve a nonsymmetric sparse matrix problem associated with a grid

     Input Parameters:
     grid - the given grid
     procinfo - the processor information (in BlockSolve format)

 +*/

void worker(par_grid *grid, BSprocinfo *procinfo)
{
	int	i, j, n, bs = 1, n_solves = 13, write_option = 0;
	BSspmat *A;
	BSpar_mat *pA, *f_pA;
	BScomm *Acomm, *f_comm;
	FLOAT shifted_diag, shifted_diag_inc, residual[2];
	int	num_iter;
	FLOAT	*x, *rhs, t;
	extern double drand48();

	/* number grid to use in matrix assembly */
	num_grid3d(grid,procinfo);

	/* now call the routines to set up the matrix */
	A = get_mat3d(grid,procinfo);

	/* Set symmetry and storage scheme to be used */
	BSset_mat_symmetric(A,grid->symmetric);
	BSset_mat_icc_storage(A,grid->icc_storage);

	/* write out matrix */
	if(write_option) write_mat_matlab("MAT.m",A,procinfo);

	/* permute the matrix */
	pA = BSmain_perm(procinfo,A); CHKERR(0);

	/* diagonally scale the matrix */
	BSscale_diag(pA,pA->diag,procinfo); CHKERR(0);

	/* set up the communication structure for triangular matrix solution */
	Acomm = BSsetup_forward(pA,procinfo); CHKERR(0);

	/* shifted_diag is the initial diagonal */
	shifted_diag = .80;
	shifted_diag_inc = 1.2/(n_solves-1);

	/* get a copy of the sparse matrix */
	f_pA = BScopy_par_mat(pA); CHKERR(0);

	/* set up a communication structure for factorization */
	f_comm = BSsetup_factor(f_pA,procinfo); CHKERR(0);

	/* set diagonal to initial shifted_diag */
	BSset_diag(f_pA,shifted_diag,procinfo); CHKERR(0);

	srand48((long)(11311));
	for (j=0; j<n_solves; j++) {
		/* factor the matrix until successful */
		while (BSfactor(f_pA,f_comm,procinfo) != 0) {
			CHKERR(0);
			/* recopy the nonzeroes */
			BScopy_nz(pA,f_pA); CHKERR(0);
			/* increment the diagonal shift */
			shifted_diag += 0.1;
			BSset_diag(f_pA,shifted_diag,procinfo); CHKERR(0);
		}
		CHKERR(0);

		/* set up the rhs and the x vector */
		n = A->num_rows;
		rhs = (FLOAT *) MALLOC(sizeof(FLOAT)*bs*n);
		x = (FLOAT *) MALLOC(sizeof(FLOAT)*bs*n);
		t = A->global_num_rows;
		t = 1.0/sqrt(t);
		for (i=0; i<n; i++) {
			rhs[i] = drand48();
			rhs[i] = t;
			x[i] = 0.0;
			if(bs==2) {
				rhs[i+n] = t*i;
				x[i+n] = 0.0;
			}
		}
	
		/* write out rhs */
		if(write_option) write_vec_matlab("RHS.m",rhs,A,procinfo);

		/* solve it */
		BSctx_set_max_it(procinfo,100);
		BSctx_set_restart(procinfo,20);
		BSctx_set_guess(procinfo,TRUE);
		BSctx_set_tol(procinfo,1.0e-7);
		if(j>0) BSctx_set_pr(procinfo,FALSE);
		num_iter = BSpar_solve(pA,f_pA,Acomm,rhs,x,residual,procinfo); CHKERR(0);

		if (procinfo->my_id == 0) {
			printf("   shifted_diag(%d) = %f; num_iter(%d) = %d;\n",j+1,shifted_diag,
				j+1,num_iter);
			/*
			printf("Took %d iterations: residuals = ",num_iter);
			for (i=0; i<bs; i++)
				printf("%e ",residual[i]);
			printf("\n");
			*/
		}

		/* write out ans */
		if(write_option) write_vec_matlab("ANS.m",x,A,procinfo);

		FREE(rhs);
		FREE(x);

		/* set up new system, increment the diagonal shift */
		BScopy_nz(pA,f_pA); CHKERR(0);
		shifted_diag += shifted_diag_inc;
		BSset_diag(f_pA,shifted_diag,procinfo); CHKERR(0);
	}

	/* free the grid */
	free_grid(grid);

	/* free the spmat */
	BSfree_easymat(A);

	/* free the par mat, etc. */
	BSfree_par_mat(pA);
	BSfree_copy_par_mat(f_pA);
	BSfree_comm(Acomm);
	BSfree_comm(f_comm);
}