#include <petsc.h>

typedef struct {
  PetscReal u, v;
} Field; /*structure for unknowns*/

typedef struct {
} AppCtx;

PetscErrorCode SetVariableBounds(DM da, Vec xl, Vec xu);
PetscErrorCode FormFunctionLocal(DMDALocalInfo *infopt, Field ***statenext, Field ***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=1, Ny=5, Nz=5;
  
  int dof=2, width=1;

  DM da;
  ierr = DMDACreate3d(PETSC_COMM_WORLD,DMDA_BOUNDARY_NONE,DMDA_BOUNDARY_NONE,DMDA_BOUNDARY_NONE, DMDA_STENCIL_STAR,Nx,Ny,Nz,PETSC_DECIDE,PETSC_DECIDE, PETSC_DECIDE,dof,width,NULL,NULL,NULL,&da);CHKERRQ(ierr);

  AppCtx user;

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

  ierr = DMSetApplicationContext(da,&user);CHKERRQ(ierr);
  ierr = DMDASNESSetFunctionLocal(da,INSERT_VALUES,(PetscErrorCode (*)(DMDALocalInfo*,void*,void*,void*))FormFunctionLocal,&user); CHKERRQ(ierr);

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

  ierr = PetscObjectSetName((PetscObject) xl, "xl"); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) xu, "xu"); CHKERRQ(ierr);

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

  //VecView(xl,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);
  //VecView(xu,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

  Vec solcur; 
  ierr = DMCreateGlobalVector(da,&solcur); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) solcur, "solcur"); CHKERRQ(ierr);

  ierr = VecSet(solcur,0.0); CHKERRQ(ierr);

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

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

  ierr = SNESSolve(snes,NULL,solcur); CHKERRQ(ierr);


  Vec myx,myb;
  Mat myjac;
  KSP ksp;
  SNESGetKSP(snes,&ksp);
  KSPGetSolution(ksp,&myx);
  KSPGetRhs(ksp,&myb);
  KSPGetOperators(ksp,&myjac,NULL,NULL); 
  MatView(myjac,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  VecView(myb,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);
  VecView(myx,PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);


  Vec gva; 
  double gna,gna2;
  ierr = SNESGetFunctionNorm(snes,&gna);
  ierr = SNESGetFunction(snes,&gva,NULL,NULL);
  VecNorm(gva,NORM_2,&gna2);

  double fna, xna;
  Vec frescur;
  VecDuplicate(solcur,&frescur);
  PetscObjectSetName((PetscObject) frescur, "frescur"); CHKERRQ(ierr); 
  
  SNESComputeFunction(snes,solcur,frescur);  

  VecNorm(frescur,NORM_2,&fna);
  VecNorm(solcur,NORM_2,&xna);
  PetscPrintf(PETSC_COMM_WORLD,"the norm by getnorm is %.16e and by getfun+vecnorm is %.16e and by computefun is %.16e\n",gna,gna2,fna);

  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;
  PetscInt mxstart, mystart, mzstart, mxlen, mylen, mzlen;
  PetscInt ix, iy, iz;
  
  PetscErrorCode ierr;

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

  ierr=DMDAGetCorners(da,&mxstart,&mystart,&mzstart,&mxlen,&mylen,&mzlen);

   for (iz=mzstart; iz<mzstart+mzlen; iz++){
    for (iy=mystart; iy<mystart+mylen; iy++) {
      for (ix=mxstart; ix<mxstart+mxlen; ix++)
	{	  
	  l[iz][iy][ix][0]=-SNES_VI_INF;
	  l[iz][iy][ix][1]=0; //-SNES_VI_INF
	  
	  u[iz][iy][ix][0]=SNES_VI_INF;
	  u[iz][iy][ix][1]=SNES_VI_INF; 
	}}}

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

   PetscFunctionReturn(0);

}


#undef __FUNCT__ 
#define __FUNCT__ "FormFunctionLocal"
PetscErrorCode FormFunctionLocal(DMDALocalInfo *infopt, Field ***statenext, Field ***residual, void *ptr)
{
 
  PetscInt ix, iy, iz, mxstart, mystart, mzstart, mxlen, mylen, mzlen, mxend, myend, mzend;
  mxstart = infopt->xs;
  mystart = infopt->ys;
  mzstart = infopt->zs;
  mxlen = infopt->xm;
  mylen = infopt->ym;
  mzlen = infopt->zm;
  mxend = mxstart + mxlen;
  myend = mystart + mylen;
  mzend = mzstart + mzlen;

  for (iz=mzstart; iz<mzend; iz++){
    for (iy=mystart; iy<myend; iy++) {
      for (ix=mxstart; ix<mxend; ix++)
	{
	  residual[iz][iy][ix].u = statenext[iz][iy][ix].u;
	  residual[iz][iy][ix].v = statenext[iz][iy][ix].v;
	}}}

  residual[0][0][0].u += 100;
  residual[0][4][0].u += 100;
  residual[0][0][4].u += 100;
  residual[0][4][4].u += 100;
  residual[0][2][2].v += 100;    

  PetscFunctionReturn(0);
}