static char help[] = "Solves 2D inhomogeneous Laplacian using multigrid.\n\n";

#include <petscdmda.h>
#include <petscksp.h>
#include <petsctime.h>

typedef struct {
    PassiveReal mu;         // viscosity
    Vec         phi;       // level set function
} Stokesparas;

typedef struct {
    PetscScalar u,v,p;
} Field;

typedef struct {
    KSP         ksp;        // linear Krylov solver
    Stokesparas para;
} Stokes2D;

extern PetscErrorCode Stokes2D_initial(Stokes2D*, DM);
extern PetscErrorCode Stokes2D_solve(Stokes2D*);
extern PetscErrorCode Stokes2D_finalize(Stokes2D*);
extern PetscErrorCode printSol(Stokes2D*);
extern PetscErrorCode ComputeMatrix(KSP,Mat,Mat,MatStructure*,void*);
extern PetscErrorCode ComputeRHS(KSP,Vec,void*);

int main(int argc,char **argv)
{
    Stokes2D       stokes;
    DM             da;
    PetscErrorCode ierr;
    PetscReal       t0,t1,t2,t3;
    
    ierr = PetscTime(&t0);CHKERRQ(ierr);
    PetscInitialize(&argc,&argv,(char*)0,help);
    ierr = DMDACreate2d(PETSC_COMM_WORLD, DMDA_BOUNDARY_NONE, DMDA_BOUNDARY_NONE,DMDA_STENCIL_STAR,-3,-3,PETSC_DECIDE,PETSC_DECIDE,3,3,0,0,&da);CHKERRQ(ierr);
    Stokes2D_initial(&stokes, da);
    ierr = PetscTime(&t1);CHKERRQ(ierr);
    Stokes2D_solve(&stokes);
    ierr = PetscTime(&t2);CHKERRQ(ierr);
    printSol(&stokes);
    ierr = PetscTime(&t3);CHKERRQ(ierr);
    
    ierr = PetscPrintf(PETSC_COMM_WORLD,"Time cost for creating solver context %g s, and for solving %g s, and for printing %g s.\n",t1-t0,t2-t1,t3-t2);CHKERRQ(ierr);
    ierr = Stokes2D_finalize(&stokes);CHKERRQ(ierr);
    ierr = PetscFinalize();CHKERRQ(ierr);
    
    return 0;
}


PetscErrorCode Stokes2D_initial(Stokes2D *s, DM da)
{
    PetscErrorCode ierr;
    ierr = KSPCreate(PETSC_COMM_WORLD,&(s->ksp));CHKERRQ(ierr);
    ierr = KSPSetComputeRHS(s->ksp,ComputeRHS,NULL);CHKERRQ(ierr);
    ierr = KSPSetComputeOperators(s->ksp,ComputeMatrix,NULL);CHKERRQ(ierr);
    ierr = KSPSetDM(s->ksp,da);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode Stokes2D_solve(Stokes2D *s)
{
    PetscErrorCode ierr;
    ierr = KSPSetFromOptions(s->ksp);CHKERRQ(ierr);
    ierr = KSPSetUp(s->ksp);CHKERRQ(ierr);
    ierr = KSPSolve(s->ksp,NULL,NULL);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode Stokes2D_finalize(Stokes2D *s)
{
    PetscErrorCode ierr;
//    ierr = DMDestroy(ka->da);CHKERRQ(ierr);
    ierr = KSPDestroy(&(s->ksp));CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode printSol(Stokes2D* s)
{
    PetscErrorCode ierr;
    PetscViewer    viewer;
    Vec            b,x;
    
    PetscFunctionBeginUser;
    ierr = KSPGetSolution(s->ksp,&x);CHKERRQ(ierr);
    ierr = KSPGetRhs(s->ksp,&b);CHKERRQ(ierr);
    ierr = PetscViewerASCIIOpen(PETSC_COMM_WORLD, "output/ksp_x.txt", &viewer);CHKERRQ(ierr);
    ierr = VecView(x, viewer);CHKERRQ(ierr);
    ierr = PetscViewerDestroy(&viewer);CHKERRQ(ierr);
    ierr = PetscViewerASCIIOpen(PETSC_COMM_WORLD, "output/ksp_b.txt", &viewer);CHKERRQ(ierr);
    ierr = VecView(b, viewer);CHKERRQ(ierr);
    ierr = PetscViewerDestroy(&viewer);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode ComputeRHS(KSP ksp,Vec b,void *ctx)
{
    PetscErrorCode ierr;
    PetscInt       i,j,mx,my,xm,ym,xs,ys;
    PetscScalar    hx,hy;
    Field          **array;
    DM             da;
    
    PetscFunctionBeginUser;
    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.0 / (PetscReal)(mx-1);
    hy   = 1.0 / (PetscReal)(my-1);
    ierr = DMDAGetCorners(da,&xs,&ys,0,&xm,&ym,0);CHKERRQ(ierr);
    ierr = DMDAVecGetArray(da, b, &array);CHKERRQ(ierr);
    for (j=ys; j<ys+ym; j++) {
        for (i=xs; i<xs+xm; i++) {
            if (i==0 || j==0 || i==mx-1 || j==my-1) {
                array[j][i].u = i*hx+j*hy;
                array[j][i].v = i*hx-j*hy;
                array[j][i].p = i*hx*j*hy;
            } else {
                array[j][i].u = 0;
                array[j][i].v = 0;
                array[j][i].p = 0;
            }
        }
    }
    ierr = DMDAVecRestoreArray(da, b, &array);CHKERRQ(ierr);
    ierr = VecAssemblyBegin(b);CHKERRQ(ierr);
    ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode ComputeMatrix(KSP ksp,Mat J,Mat jac,MatStructure *str,void *ctx)
{
    PetscErrorCode ierr;
    PetscInt       i,j,mx,my,xm,ym,xs,ys;
    PetscScalar    v[5];
    PetscScalar    hx, hy, dhx, dhy, dhx2, dhy2, hxdhy, hydhx;
    MatStencil     row, col[5];
    DM             da;
    
    PetscFunctionBeginUser;
    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.0 / (PetscReal)(mx-1);
    hy        = 1.0 / (PetscReal)(my-1);
    dhx       = 1.0/hx;
    dhy       = 1.0/hy;
    dhx2      = 1.0/hx/hx;
    dhy2      = 1.0/hy/hy;
    hxdhy     = hx/hy;
    hydhx     = hy/hx;
    ierr      = DMDAGetCorners(da,&xs,&ys,0,&xm,&ym,0);CHKERRQ(ierr);
    for (j=ys; j<ys+ym; j++) {
        for (i=xs; i<xs+xm; i++) {
            row.i = i; row.j = j;  row.c = 0;
            if (i==0 || j==0 || i==mx-1 || j==my-1) {
                v[0]=1;   row.c = 0;  ierr = MatSetValuesStencil(jac,1,&row,1,&row,v,INSERT_VALUES);CHKERRQ(ierr);
                v[0]=1;   row.c = 1;  ierr = MatSetValuesStencil(jac,1,&row,1,&row,v,INSERT_VALUES);CHKERRQ(ierr);
                v[0]=1;   row.c = 2;  ierr = MatSetValuesStencil(jac,1,&row,1,&row,v,INSERT_VALUES);CHKERRQ(ierr);
                
//                ierr = PetscPrintf(PETSC_COMM_WORLD,"1. Creating matrix %d %d %d %d %d %d.\n",row.i,row.j,row.c,col[0].i,col[0].j,col[0].c);CHKERRQ(ierr);
            } else {
                row.c = 0;
                v[0] = dhx2;            col[0].i = i;   col[0].j = j-1; col[0].c = 0;
                v[1] = dhx2;            col[1].i = i-1; col[1].j = j;   col[1].c = 0;
                v[2] = -2*dhx2-2*dhy2;  col[2].i = i;   col[2].j = j;   col[2].c = 0;
                v[3] = dhy2;            col[3].i = i;   col[3].j = j+1; col[3].c = 0;
                v[4] = dhy2;            col[4].i = i+1; col[4].j = j;   col[4].c = 0;
                ierr = MatSetValuesStencil(jac,1,&row,5,col,v,INSERT_VALUES);CHKERRQ(ierr);
                row.c = 1;
                v[0] = dhx2;            col[0].i = i;   col[0].j = j-1; col[0].c = 1;
                v[1] = dhx2;            col[1].i = i-1; col[1].j = j;   col[1].c = 1;
                v[2] = -2*dhx2-2*dhy2;  col[2].i = i;   col[2].j = j;   col[2].c = 1;
                v[3] = dhy2;            col[3].i = i;   col[3].j = j+1; col[3].c = 1;
                v[4] = dhy2;            col[4].i = i+1; col[4].j = j;   col[4].c = 1;
                ierr = MatSetValuesStencil(jac,1,&row,5,col,v,INSERT_VALUES);CHKERRQ(ierr);
                row.c = 2;
                v[0] = dhx2;            col[0].i = i;   col[0].j = j-1; col[0].c = 2;
                v[1] = dhx2;            col[1].i = i-1; col[1].j = j;   col[1].c = 2;
                v[2] = -2*dhx2-2*dhy2;  col[2].i = i;   col[2].j = j;   col[2].c = 2;
                v[3] = dhy2;            col[3].i = i;   col[3].j = j+1; col[3].c = 2;
                v[4] = dhy2;            col[4].i = i+1; col[4].j = j;   col[4].c = 2;
                ierr = MatSetValuesStencil(jac,1,&row,5,col,v,INSERT_VALUES);CHKERRQ(ierr);
//                ierr = PetscPrintf(PETSC_COMM_WORLD,"2. Creating matrix %d %d %d %d %d %d.\n",row.i,row.j,row.c,col[0].i,col[0].j,col[0].c);CHKERRQ(ierr);
            }
        }
    }
    ierr = MatAssemblyBegin(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    ierr = MatAssemblyEnd(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}