#include "petscsolver.h"

PETScSolver::PETScSolver(int argc,char **argv,int g_Ni,int g_Nj,int g_Nk,int *dims,int *l_Ni,int *l_Nj,int *l_Nk,int _Nall,double h)
{
  MPI_Comm_rank(MPI_COMM_WORLD,&rank);

	Nall = _Nall;
	ih = 1./h;
	iepsilon0 = 1./8.854187817E-12;
	
	bound = NULL;

	PetscInitialize(&argc,&argv,(char*)0,(char*)0);

	KSPCreate(PETSC_COMM_WORLD,&ksp);
	DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,
      DMDA_STENCIL_BOX,g_Ni,g_Nj,g_Nk,dims[0],dims[1],dims[2],1,1,l_Ni,l_Nj,l_Nk,&da);
	KSPSetDM(ksp,da);

  DMCreateGlobalVector(da,&b);
  VecDuplicate(b,&x);
  VecSet(x,0);
  VecSet(b,0);
}

PETScSolver::~PETScSolver()
{
	VecDestroy(&x);
	VecDestroy(&b);
	DMDestroy(&da);
	KSPDestroy(&ksp);
	PetscFinalize();
}

/* Needs to be called before initsolve() or solve() */
void PETScSolver::init(int *_bound)
{
  bound = _bound;
	KSPSetComputeOperators(ksp,ComputeMatrix,(void*)bound);
}

void PETScSolver::initsolve(double *phi,double *charge,double reltol)
{
  KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
  KSPSetFromOptions(ksp);
  KSPSetTolerances(ksp,reltol,PETSC_DEFAULT,PETSC_DEFAULT,PETSC_DEFAULT);

  PetscScalar *xarray;
  VecGetArray(x,&xarray);
  for (int i=0; i<Nall; i++)
    xarray[i] = phi[i];
  VecRestoreArray(x, &xarray);

  PetscScalar *barray;
  VecGetArray(b,&barray);
  for (int i=0; i<Nall; i++) {
    if (bound[i]==0)
      barray[i] = charge[i]*ih*iepsilon0;
    else
      barray[i] = phi[i];
  }
  VecRestoreArray(b,&barray);

  KSPSolve(ksp,b,x);

  int its;
  KSPGetIterationNumber(ksp,&its);
  if (rank==0) std::cout<<"Initsolve: "<<its<<" iterations\n";
  
	KSPConvergedReason reason;
	KSPGetConvergedReason(ksp,&reason);
	if (reason<0) {
		if (rank==0) std::cerr<<"Initsolve: no convergence\n";
		MPI_Finalize();
		exit(101);
	}

  KSPGetSolution(ksp,&x);
  VecGetArray(x,&xarray);
  for (int i=0; i<Nall; i++)
    phi[i] = xarray[i];
  VecRestoreArray(x,&xarray);
}

void PETScSolver::solve(int timestep,double *phi,double *charge,double reltol)
{
  KSPSetTolerances(ksp,reltol,PETSC_DEFAULT,PETSC_DEFAULT,PETSC_DEFAULT);

  PetscScalar *barray;
  VecGetArray(b,&barray);
  for (int i=0; i<Nall; i++) {
    if (bound[i]==0)
      barray[i] = charge[i]*ih*iepsilon0;
    else
      barray[i] = phi[i];
  }
  VecRestoreArray(b,&barray);

  KSPSolve(ksp,b,x);

  int its;
  KSPGetIterationNumber(ksp,&its);
  if (rank==0) std::cout<<"Timestep "<<timestep<<": "<<its<<" iterations\n";
  
  KSPGetSolution(ksp,&x);
  PetscScalar *xarray;
  VecGetArray(x,&xarray);
  for (int i=0; i<Nall; i++)
    phi[i] = xarray[i];
  VecRestoreArray(x,&xarray);
}

PetscErrorCode ComputeMatrix(KSP ksp,Mat A,Mat B,void *ctx)
{
	DM             da;
	PetscErrorCode ierr;
	PetscInt       i,j,k,mx,my,mz,xm,ym,zm,xs,ys,zs;
	PetscScalar    v[7];
	MatStencil     row,col[7];
	int *bound = (int *)ctx

	PetscFunctionBeginUser;
	ierr    = KSPGetDM(ksp,&da);CHKERRQ(ierr);
	ierr    = DMDAGetInfo(da,0,&mx,&my,&mz,0,0,0,0,0,0,0,0,0);CHKERRQ(ierr);
	ierr    = DMDAGetCorners(da,&xs,&ys,&zs,&xm,&ym,&zm);CHKERRQ(ierr);

	for (k=zs; k<zs+zm; k++) {
		for (j=ys; j<ys+ym; j++) {
			for (i=xs; i<xs+xm; i++) {
				row.i = i; row.j = j; row.k = k;
				if (bound[(i-xs)+(j-ys)*xm+(k-zs)*xm*ym]==1) { //Dirichlet nodes
					v[0] = 1.0;
					ierr = MatSetValuesStencil(B,1,&row,1,&row,v,INSERT_VALUES);CHKERRQ(ierr);
				} else {                                       //others
					v[0] = -1; col[0].i = i; col[0].j = j; col[0].k = k-1;
					v[1] = -1; col[1].i = i; col[1].j = j-1; col[1].k = k;
					v[2] = -1; col[2].i = i-1; col[2].j = j; col[2].k = k;
					v[3] =  6; col[3].i = row.i; col[3].j = row.j; col[3].k = row.k;
					v[4] = -1; col[4].i = i+1; col[4].j = j; col[4].k = k;
					v[5] = -1; col[5].i = i; col[5].j = j+1; col[5].k = k;
					v[6] = -1; col[6].i = i; col[6].j = j; col[6].k = k+1;
					ierr = MatSetValuesStencil(B,1,&row,7,col,v,INSERT_VALUES);CHKERRQ(ierr);
				}
			}
		}
	}

	ierr   = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
	ierr   = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
	PetscFunctionReturn(0);
}