#include <petscksp.h>
#include <petscdm.h>
#include <petscdmda.h>
#include <petsc.h>
#include <petscsnes.h>
#include "petscsys.h"
#include "petscis.h"
#include <stdio.h>

int GlobalToLocal(int x, int y)
{
    if((x == 0) && (y == 0)) return 0;
    if((x == 1) && (y == 0)) return 1;
    if((x == 0) && (y == 1)) return 2;
    if((x == 1) && (y == 1)) return 3;
}

int main(int argc, char *argv[])
{
    PetscInitialize(&argc, &argv, PETSC_NULLPTR, PETSC_NULLPTR);
    int rank, size;
    MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
    MPI_Comm_size(PETSC_COMM_WORLD, &size);
    
    double kloc[16];
    kloc[0] =  0.6666666667; kloc[1] = -0.1666666667; kloc[2] = -0.1666666667; kloc[3] = -0.3333333333;
    kloc[4] = -0.1666666667;  kloc[5] =  0.6666666667; kloc[6] = -0.3333333333; kloc[7] = -0.1666666667;
    kloc[8] = -0.1666666667; kloc[9] = -0.3333333333; kloc[10] = 0.6666666667; kloc[11] = -0.1666666667;
    kloc[12] = -0.3333333333; kloc[13] = -0.1666666667; kloc[14] = -0.1666666667; kloc[15] = 0.6666666667;

    double floc[4];
    floc[0] = 64; floc[1] = 64; floc[2] = 64; floc[3] = 64;

    int bnd[17];
    bnd[0] = 20; bnd[1] = 21; bnd[2] = 22; bnd[3] = 23; bnd[4] = 24;
    bnd[5] = 56; bnd[6] = 57; bnd[7] = 58; bnd[8] = 59; bnd[9] = 60;
    bnd[10] = 29; bnd[11] = 38; bnd[12] = 47;
    bnd[13] = 33; bnd[14] = 42; bnd[15] = 51;
    bnd[16] = 40;
    
    PetscErrorCode ierr;
    PetscInt N = 81;
    
    PetscInt d_nz = 10;
    PetscInt o_nz = 10;

    KSP       ksp;
    Vec       x, b;
    Mat       A;
    
    ierr = MatCreate(PETSC_COMM_WORLD, &A); CHKERRQ(ierr);
    ierr = MatSetType(A, MATMPIAIJ); CHKERRQ(ierr);
    ierr = MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, N, N); CHKERRQ(ierr);
    ierr = MatMPIAIJSetPreallocation(A,d_nz, NULL, o_nz, NULL); CHKERRQ(ierr);
    ierr = MatSetOption(A,MAT_SYMMETRIC,PETSC_TRUE); CHKERRQ(ierr);
    ierr = MatCreateVecs(A, &b, &x);CHKERRQ(ierr);

    
    int ncells = 8;
    int nodes[4];
    
    PetscInt i, j, nx, ny;
    
    for (j = 2; j < 6; j++)
      {
	for (i = 2; i < 6; i++)
	  {
	    for (ny = 0; ny < 2; ny++)
	      {
		for (nx = 0; nx < 2; nx++)
		  {
		    int ni, nj, gobaln;
		 
		    ni = i + nx;
		    nj = j + ny;
		    
		    // one index is used to refer a global node
		    gobaln = ni * (ncells + 1) + nj;

		    //printf("%d\t%d\t%d\n", i, j, gobaln);
		   
		    nodes[GlobalToLocal(nx, ny)] = gobaln;
		  }
	      }
	    
	    MatSetValues(A,  4, nodes,  4, nodes, kloc, ADD_VALUES);
	    VecSetValues(b,  4, nodes, floc, ADD_VALUES);
	  }
      }

    ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    
    ierr = VecAssemblyBegin(b); CHKERRQ(ierr);
    ierr = VecAssemblyEnd(b); CHKERRQ(ierr);

    ierr =  MatZeroRowsColumns(A, 17, bnd, 1, x, b); CHKERRQ(ierr);


    KSPCreate(PETSC_COMM_WORLD, &ksp);
    KSPSetOperators(ksp, A, A);
    
    ierr = PetscOptionsSetValue(NULL,"-ksp_type", "preonly"); CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL,"-pc_type", "redistribute"); CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL,"-redistribute_ksp_type", "cg"); CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL,"-redistribute_pc_type", "jacobi"); CHKERRQ(ierr);
    //ierr = PetscOptionsSetValue(NULL, "-ksp_monitor", NULL);CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL, "-redistribute_ksp_view", NULL);CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL, "-redistribute_ksp_monitor_true_residual", NULL);CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL, "-redistribute_ksp_converged_reason", NULL);CHKERRQ(ierr);
    ierr = PetscOptionsSetValue(NULL, "-options_left", NULL);CHKERRQ(ierr);
    ierr = KSPSetFromOptions(ksp); CHKERRQ(ierr);
    
    KSPSolve(ksp, b, x);
    
    PetscFinalize();
    
   return 0;
}