#include <petscdmstag.h>

#define ELEMENT DMSTAG_ELEMENT

int main(int argc, char **argv)
{
    PetscCall(PetscInitialize(&argc, &argv, (char *)0, NULL));
    PetscInt M = 4, N = 4;
    PetscScalar H_x = 1.0 / (double)M, H_y = 1.0 / (double)N, avg_kappa_e;
    DM dm_P;
    PetscCall(DMStagCreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, M, N, PETSC_DECIDE, PETSC_DECIDE, 0, 0, 1, DMSTAG_STENCIL_BOX, 1, NULL, NULL, &dm_P));
    // Only one DOF in the center of element.
    PetscCall(DMSetUp(dm_P));
    Mat A;
    PetscCall(DMCreateMatrix(dm_P, &A));
    PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz, extrax, extray, extraz;
    PetscCall(DMStagGetCorners(dm_P, &startx, &starty, &startz, &nx, &ny, &nz, &extrax, &extray, &extraz));
    PetscScalar meas_elem = H_x * H_y;
    for (ey = starty; ey < starty + ny; ++ey)
        for (ex = startx; ex < startx + nx; ++ex)
        {
            if (ex >= 1)
            {
                DMStagStencil row[2], col[2];
                PetscScalar val_A[2][2];
                row[0] = (DMStagStencil){.i = ex - 1, .j = ey, .c = 0, .loc = ELEMENT};
                row[1] = (DMStagStencil){.i = ex, .j = ey, .c = 0, .loc = ELEMENT};
                col[0] = (DMStagStencil){.i = ex - 1, .j = ey, .c = 0, .loc = ELEMENT};
                col[1] = (DMStagStencil){.i = ex, .j = ey, .c = 0, .loc = ELEMENT};
                avg_kappa_e = 1.0;
                val_A[0][0] = H_y * H_y / meas_elem * avg_kappa_e;
                val_A[0][1] = -H_y * H_y / meas_elem * avg_kappa_e;
                val_A[1][0] = -H_y * H_y / meas_elem * avg_kappa_e;
                val_A[1][1] = H_y * H_y / meas_elem * avg_kappa_e;
                PetscCall(DMStagMatSetValuesStencil(dm_P, A, 2, &row[0], 2, &col[0], &val_A[0][0], ADD_VALUES));
            }
            if (ey >= 1)
            {
                DMStagStencil row[2], col[2];
                PetscScalar val_A[2][2];
                row[0] = (DMStagStencil){.i = ex, .j = ey - 1, .c = 0, .loc = ELEMENT};
                row[1] = (DMStagStencil){.i = ex, .j = ey, .c = 0, .loc = ELEMENT};
                col[0] = (DMStagStencil){.i = ex, .j = ey - 1, .c = 0, .loc = ELEMENT};
                col[1] = (DMStagStencil){.i = ex, .j = ey, .c = 0, .loc = ELEMENT};
                avg_kappa_e = 2.0;
                val_A[0][0] = H_x * H_x / meas_elem * avg_kappa_e;
                val_A[0][1] = -H_x * H_x / meas_elem * avg_kappa_e;
                val_A[1][0] = H_x * H_x / meas_elem * avg_kappa_e;
                val_A[1][1] = H_x * H_x / meas_elem * avg_kappa_e;
                PetscCall(DMStagMatSetValuesStencil(dm_P, A, 2, &row[0], 2, &col[0], &val_A[0][0], ADD_VALUES));
            }
        }
    PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Successfully construct the Mat!\n"));
    PetscCall(MatDestroy(&A));
    PetscCall(DMDestroy(&dm_P));

    return 0;
}