#include <petsc.h>

PetscErrorCode SetVariableBounds(DM da, Vec xl, Vec xu);
PetscErrorCode FormFunctionLocal(DMDALocalInfo *infopt, double *statenext, double *residual, void *ptr);

#undef __FUNCT__
#define __FUNCT__ "main"

int main(int argc, char **argv)
{
  PetscInitialize(&argc, &argv, PETSC_NULL, PETSC_NULL);
  PetscPrintf(PETSC_COMM_WORLD,"----Initializing------\n");
  PetscErrorCode ierr;

  int myrank;
  MPI_Comm_rank(MPI_COMM_WORLD,&myrank);

  PetscInt Nx=4;
  
  int dof=1, width=1;

  DM da;
  ierr = DMDACreate1d(PETSC_COMM_WORLD,DMDA_BOUNDARY_NONE,Nx,dof,width,NULL,&da);CHKERRQ(ierr);

  SNES snes;
  ierr = SNESCreate(PETSC_COMM_WORLD,&snes); CHKERRQ(ierr);
  ierr = SNESSetDM(snes, (DM) da); CHKERRQ(ierr);

  ierr = DMDASNESSetFunctionLocal(da,INSERT_VALUES,(PetscErrorCode (*)(DMDALocalInfo*,void*,void*,void*))FormFunctionLocal,NULL); CHKERRQ(ierr);

  /* set lower and upper bounds for solutions */
  Vec xl, xu;
  ierr = DMCreateGlobalVector(da,&xl); CHKERRQ(ierr);
  ierr = DMCreateGlobalVector(da,&xu); CHKERRQ(ierr);

  SetVariableBounds(da,xl,xu);
  ierr = SNESVISetVariableBounds(snes,xl,xu); CHKERRQ(ierr);

  Vec solcur; 
  ierr = DMCreateGlobalVector(da,&solcur); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) solcur, "solcur"); CHKERRQ(ierr);
  
  /* set the initial guess */
  ierr = VecSet(solcur,0.0); CHKERRQ(ierr);

  ierr = SNESSetFromOptions(snes); CHKERRQ(ierr);

  VecView(solcur,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

  /* solve the equation */
  ierr = SNESSolve(snes,NULL,solcur); CHKERRQ(ierr);

  /* view the kspmat and ksprhs */
  Vec myx,myb;
  Mat myjac;
  KSP ksp;
  SNESGetKSP(snes,&ksp);
  KSPGetSolution(ksp,&myx);
  KSPGetRhs(ksp,&myb);
  KSPGetOperators(ksp,&myjac,NULL,NULL); 

  ierr = PetscObjectSetName((PetscObject) myx, "myx"); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) myb, "myb"); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) myjac, "myjac"); CHKERRQ(ierr);

  MatView(myjac,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  VecView(myb,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);
  VecView(myx,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

  VecView(solcur,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

  /*--------finalzie and exit the program --------*/
  ierr = PetscFinalize(); CHKERRQ(ierr);
  PetscFunctionReturn(0);
}


#undef __FUNCT__
#define __FUNCT__ "SetVariableBounds"
PetscErrorCode SetVariableBounds(DM da, Vec xl, Vec xu)
{
  PetscScalar *l, *u;
  PetscErrorCode ierr;

  ierr=DMDAVecGetArray(da,xl,&l); CHKERRQ(ierr);
  ierr=DMDAVecGetArray(da,xu,&u); CHKERRQ(ierr);

  double ulbd, uubd, vlbd, vubd;
  ulbd=-SNES_VI_INF;
  uubd=SNES_VI_INF;
  
  vlbd = 0; //-SNES_VI_INF;
  vubd = SNES_VI_INF;
     	  
  l[0] = ulbd;
  l[1] = vlbd;
  l[2] = ulbd;
  l[3] = vlbd;

  u[0] = uubd;
  u[1] = vubd;
  u[2] = uubd;
  u[3] = vubd;

  ierr=DMDAVecRestoreArray(da,xl,&l); CHKERRQ(ierr);
  ierr=DMDAVecRestoreArray(da,xu,&u); CHKERRQ(ierr);
  
  PetscFunctionReturn(0);

}


#undef __FUNCT__ 
#define __FUNCT__ "FormFunctionLocal"
PetscErrorCode FormFunctionLocal(DMDALocalInfo *infopt, double *statenext, double *residual, void *ptr)
{
  residual[0] = statenext[0] + statenext[1];
  residual[1] = statenext[0] - statenext[1];
  residual[2] = statenext[2] + statenext[3];
  residual[3] = statenext[2] - statenext[3];

  residual[0] += 100;
  residual[3] += 50;
 
  PetscFunctionReturn(0);
}