#include "pargrid.h"

/*+ worker0 - Solve a 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, global_nnz, num_iter, write_option = 0;
	BSspmat *A;
	BSpar_mat *pA, *f_pA;
	BScomm *Acomm, *f_comm;
	FLOAT shifted_diag, residual;
	FLOAT	*x, *rhs, t;

	/* 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);

	/* count nnzs for display */
	global_nnz = 2*pA->local_nnz - pA->num_rows;
	GISUM(&global_nnz,1,&i,procinfo->procset);
	if(procinfo->my_id==0) {
		printf("o  ");
		printf("Number of nonzeros = %d\n",global_nnz);
	}

	/* 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);

	/* now factor the matrix */
	/* set the initial diagonal to 1.0 */
	shifted_diag = 1.0;
	/* 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);
	/* 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 */
	rhs = (FLOAT *) MALLOC(sizeof(FLOAT)*A->num_rows);
	x = (FLOAT *) MALLOC(sizeof(FLOAT)*A->num_rows);
	t = A->global_num_rows;
	t = 1.0/sqrt(t);
	for (i=0;i<A->num_rows;i++) {
		rhs[i] = t;
		x[i] = 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_method(procinfo,CG);
	BSctx_set_tol(procinfo,1.0e-5);
	num_iter = BSpar_solve(pA,f_pA,Acomm,rhs,x,&residual,procinfo); CHKERR(0);

	if (procinfo->my_id == 0) {
		printf("o  Took %d iterations: residual = %e\n",num_iter,residual);
	}

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

	FREE(rhs);
	FREE(x);

	/* 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);
}