/*
 * ex193.c
 *
 *  Created on: Jul 29, 2015
 *      Author: Fande Kong fdkong.jd@gmail.com
 */
/*
 * An example demonstrates how to use hierarchical partitioning approach
 */

#include <petscmat.h>

static char help[] = "Illustrates use of hierarchical partitioning.\n";

int main(int argc,char **args)
{
  Mat             A;                      /* matrix */
  PetscInt        m,n;                    /* mesh dimensions in x- and y- directions */
  PetscInt        i,j,Ii,J,Istart,Iend;
  PetscErrorCode  ierr;
  PetscMPIInt     size;
  PetscScalar     v;
  MatPartitioning part;
  IS              coarseparts,fineparts;
  IS              is,isn,isrows;
  MPI_Comm        comm;

  PetscCall(PetscInitialize(&argc,&args,(char*)0,help));
  comm = PETSC_COMM_WORLD;
  PetscCallMPI(MPI_Comm_size(comm,&size));
  ierr = PetscOptionsBegin(comm,NULL,"ex193","hierarchical partitioning");PetscCall(ierr);
  m = 15;
  PetscCall(PetscOptionsInt("-M","Number of mesh points in the x-direction","partitioning",m,&m,NULL));
  n = 15;
  PetscCall(PetscOptionsInt("-N","Number of mesh points in the y-direction","partitioning",n,&n,NULL));
  ierr = PetscOptionsEnd();PetscCall(ierr);

  /*
     Assemble the matrix for the five point stencil (finite difference), YET AGAIN
  */
  PetscCall(MatCreate(comm,&A));
  PetscCall(MatSetSizes(A,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n));
  PetscCall(MatSetFromOptions(A));
  PetscCall(MatSetUp(A));
  PetscCall(MatGetOwnershipRange(A,&Istart,&Iend));
  for (Ii=Istart; Ii<Iend; Ii++) {
    v = -1.0; i = Ii/n; j = Ii - i*n;
    if (i>0)   {J = Ii - n; PetscCall(MatSetValues(A,1,&Ii,1,&J,&v,INSERT_VALUES));}
    if (i<m-1) {J = Ii + n; PetscCall(MatSetValues(A,1,&Ii,1,&J,&v,INSERT_VALUES));}
    if (j>0)   {J = Ii - 1; PetscCall(MatSetValues(A,1,&Ii,1,&J,&v,INSERT_VALUES));}
    if (j<n-1) {J = Ii + 1; PetscCall(MatSetValues(A,1,&Ii,1,&J,&v,INSERT_VALUES));}
    v = 4.0; PetscCall(MatSetValues(A,1,&Ii,1,&Ii,&v,INSERT_VALUES));
  }
  PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
  PetscCall(MatView(A,PETSC_VIEWER_STDOUT_WORLD));
  /*
   Partition the graph of the matrix
  */
  PetscCall(MatPartitioningCreate(comm,&part));
  PetscCall(MatPartitioningSetAdjacency(part,A));
  PetscCall(MatPartitioningSetType(part,MATPARTITIONINGHIERARCH));
  PetscCall(MatPartitioningHierarchicalSetNcoarseparts(part,2));
  PetscCall(MatPartitioningHierarchicalSetNfineparts(part,4));
  PetscCall(MatPartitioningSetFromOptions(part));
  /* get new processor owner number of each vertex */
  PetscCall(MatPartitioningApply(part,&is));
  /* coarse parts */
  PetscCall(MatPartitioningHierarchicalGetCoarseparts(part,&coarseparts));
  PetscCall(ISView(coarseparts,PETSC_VIEWER_STDOUT_WORLD));
  /* fine parts */
  PetscCall(MatPartitioningHierarchicalGetFineparts(part,&fineparts));
  PetscCall(ISView(fineparts,PETSC_VIEWER_STDOUT_WORLD));
  /* partitioning */
  PetscCall(ISView(is,PETSC_VIEWER_STDOUT_WORLD));
  /* get new global number of each old global number */
  PetscCall(ISPartitioningToNumbering(is,&isn));
  PetscCall(ISView(isn,PETSC_VIEWER_STDOUT_WORLD));
  PetscCall(ISBuildTwoSided(is,NULL,&isrows));
  PetscCall(ISView(isrows,PETSC_VIEWER_STDOUT_WORLD));
  PetscCall(ISDestroy(&is));
  PetscCall(ISDestroy(&coarseparts));
  PetscCall(ISDestroy(&fineparts));
  PetscCall(ISDestroy(&isrows));
  PetscCall(ISDestroy(&isn));
  PetscCall(MatPartitioningDestroy(&part));
  PetscCall(MatDestroy(&A));
  PetscCall(PetscFinalize());
  return 0;
}

/*TEST

   test:
      nsize: 4
      args: -mat_partitioning_hierarchical_Nfineparts 2
      requires: parmetis
      TODO: cannot run because parmetis does reproduce across all machines, probably due to nonportable random number generator

TEST*/