static char help[] = "Parallel vector layout.\n\n";

#include <petscsys.h>
#include <petscvec.h>
#include <petscdmda.h>
#include <petscksp.h>

DM da_sc;
KSP ksp;

PetscErrorCode ierr;

struct UserContext {
  Vec rhs;
};

int N = 10;

#undef __FUNCT__
#define __FUNCT__ "ComputeRHS"
PetscErrorCode ComputeRHS(KSP ksp, Vec b, void *ctx) {
  
  UserContext*  user = (UserContext*) ctx;
  
  DM da;
  PetscInt       mx, my;
  PetscScalar    Hx, Hy;
  
  ierr = KSPGetDM(ksp,&da);CHKERRQ(ierr);
  ierr = DMDAGetInfo(da, 0, &mx, &my, 0,0,0,0,0,0,0,0,0,0);CHKERRQ(ierr);
  Hx = 1. / (PetscReal) (mx);
  Hy = 1. / (PetscReal) (my);
  
  ierr = VecScale(user->rhs, Hx * Hy); CHKERRQ(ierr);
  ierr = VecCopy(user->rhs, b); CHKERRQ(ierr);
  
  MatNullSpace nullspace;
  ierr = MatNullSpaceCreate(PETSC_COMM_WORLD,PETSC_TRUE,0,0,&nullspace); CHKERRQ(ierr);
  ierr = MatNullSpaceRemove(nullspace, b); CHKERRQ(ierr);
  ierr = MatNullSpaceDestroy(&nullspace); CHKERRQ(ierr);
  
  ierr = VecAssemblyBegin(b);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
    
  return 0;
}


#undef __FUNCT__
#define __FUNCT__ "ComputeMatrix"
PetscErrorCode ComputeMatrix(KSP ksp, Mat J, Mat jac, void *ctx) {
  
  PetscInt       mx, my, xm, ym, xs, ys;
  PetscScalar    Hx, Hy;
  DM             da;
  MatStencil          row;
  PetscScalar         HydHx, HxdHy;
  PetscScalar         v_c;
  PetscInt            ncols;
  
  ierr  = KSPGetDM(ksp, &da);CHKERRQ(ierr);
  ierr  = DMDAGetInfo(da, 0, &mx, &my, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); CHKERRQ(ierr);
  Hx = 1. / (PetscReal) (mx);
  Hy = 1. / (PetscReal) (my);
    

  ierr = MatSetOption(jac, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE); CHKERRQ(ierr);
//     ierr = MatSetOption(jac, MAT_SYMMETRIC, PETSC_TRUE); CHKERRQ(ierr);
  
  HxdHy = Hx/Hy;
  HydHx = Hy/Hx;
  ierr  = DMDAGetCorners(da, &xs, &ys, 0, &xm, &ym, 0); CHKERRQ(ierr);
  PetscScalar    v[5];
  int            k = 0;
  for (int j=ys; j<ys+ym; ++j) {
    for (int i=xs; i<xs+xm; ++i) {
      ncols = 0;
      MatStencil    col[5] = {{0}};
      v_c = 0;
      row.k = k;             row.j = j;          row.i = i;          row.c = 0;
      col[ncols].k=k;        col[ncols].j = j;   col[ncols].i = i-1; v[ncols++] = HydHx; v_c -= HydHx;
      col[ncols].k=k;        col[ncols].j = j;   col[ncols].i = i+1; v[ncols++] = HydHx; v_c -= HydHx;
      col[ncols].k=k;        col[ncols].j = j-1; col[ncols].i = i;   v[ncols++] = HxdHy; v_c -= HxdHy;
      col[ncols].k=k;        col[ncols].j = j+1; col[ncols].i = i;   v[ncols++] = HxdHy; v_c -= HxdHy;
      col[ncols].k=k;        col[ncols].j=j;     col[ncols].i = i;   v[ncols++] = v_c;
      ierr = MatSetValuesStencil(jac, 1, &row, ncols, col, v, INSERT_VALUES); CHKERRQ(ierr);
    }
  }
    
  ierr = MatAssemblyBegin(jac, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(jac, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  
  MatNullSpace nullspace;
  ierr = MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_TRUE, 0, 0, &nullspace);CHKERRQ(ierr);
  ierr = MatSetNullSpace(jac,nullspace); CHKERRQ(ierr);
  ierr = MatNullSpaceDestroy(&nullspace); CHKERRQ(ierr);
  
  return 0;
}

#undef __FUNCT__
#define __FUNCT__ "setup"
PetscErrorCode setup(UserContext * ctx) {
  
  ierr = KSPCreate(PETSC_COMM_WORLD, &ksp); CHKERRQ(ierr);
  
  ierr = KSPSetDMActive(ksp, PETSC_FALSE); CHKERRQ(ierr);
  ierr = KSPSetDM(ksp, da_sc); CHKERRQ(ierr);
  ierr = KSPSetComputeOperators(ksp, ComputeMatrix, ctx); CHKERRQ(ierr);
  ierr = KSPSetComputeRHS(ksp, ComputeRHS, ctx); CHKERRQ(ierr);
  
  ierr = KSPSetFromOptions(ksp); CHKERRQ(ierr);
  
  ierr = KSPSetUp(ksp); CHKERRQ(ierr);
  
  return 0;
}

int main(int argc, char **argv) {
  
  PetscInitialize(&argc, &argv, (char*)0, help);
  
  UserContext ctx;
  Vec phi;
  PetscInt       mx, my;
  int rank;
  
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  
  ++rank;
  
  ierr = DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR,
                      128, 128, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, NULL, &da_sc); CHKERRQ(ierr);
  ierr = DMDASetUniformCoordinates(da_sc, 0., 1., 0., 1., 0., 1.); CHKERRQ(ierr);
  ierr = DMCreateGlobalVector(da_sc, &ctx.rhs); CHKERRQ(ierr);
  ierr  = DMDAGetInfo(da_sc, 0, &mx, &my, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); CHKERRQ(ierr);
  
  ierr = KSPCreate(PETSC_COMM_WORLD, &ksp); CHKERRQ(ierr);
  ierr = VecSet(ctx.rhs, 1); CHKERRQ(ierr);
  
  ierr = setup(&ctx); CHKERRQ(ierr);
  
  PetscViewer viewer;
  PetscViewerDrawOpen(PETSC_COMM_WORLD,NULL,NULL,0,0,300,300,&viewer);
  for(int i = 0; i < N; ++i) {
    
    ierr = VecSet(ctx.rhs, (float) i * mx); CHKERRQ(ierr);
    ierr = VecScale(ctx.rhs,  rank); CHKERRQ(ierr);
    ierr = KSPSolve(ksp, NULL, NULL); CHKERRQ(ierr);
    ierr = KSPGetSolution(ksp, &phi); CHKERRQ(ierr);
    ierr = VecView(ctx.rhs, viewer); CHKERRQ(ierr);
    ierr = VecView(phi, viewer); CHKERRQ(ierr);
  }
  PetscViewerDestroy(&viewer);
  VecDestroy(&phi);
  
  PetscFinalize();
  return 0;
}