// DMs for grid hierachy
DM *da_list, *daclist;
Mat R;

ierr = PetscMalloc(sizeof(DM) * nlvls, &da_list); CHKERRQ(ierr);
ierr = PetscMalloc(sizeof(DM) * nlvls, &daclist); CHKERRQ(ierr);

// Set 0 to the finest level
daclist[0] = da_nodal;
// Set up the coarse meshes
ierr = DMCoarsenHierarchy(da_nodal, nlvls - 1, &daclist[1]); CHKERRQ(ierr);
for (PetscInt k = 0; k < nlvls; k++) {
	da_list[k] = daclist[nlvls - 1 - k];
}

ierr = PCMGSetLevels(ipc, nlvls, NULL); CHKERRQ(ierr);
ierr = PCMGSetType(ipc, PC_MG_MULTIPLICATIVE); CHKERRQ(ierr); 
ierr = PCMGSetCycleType(ipc, PC_MG_CYCLE_V); CHKERRQ(ierr);
ierr = PCMGSetGalerkin(ipc, PC_MG_GALERKIN_BOTH); CHKERRQ(ierr);

// NEW CODE
// Temporary flag to be able to test pcmg with and without redistribute
PetscBool pcmgredistribute = PETSC_FALSE;
PetscOptionsGetBool(NULL, NULL, "-pcmgredistribute", &pcmgredistribute, NULL);
IS isfreef, isfreec; // free DOFs on fine and coarse grid
Mat Rred; 	// Reduced-size interpolation matrix. In Matlab-notation: Rred = R(isfreef,isfreec)						
VecTagger tagger;
VecTaggerBox taggerbox;
Vec NNf, NNc;  // NNf[i] = 1 if DOF i is free. NNf[i] = 0 if it is locked
if (pcmgredistribute==PETSC_TRUE)
{
	ierr = DMCreateGlobalVector(da_list[nlvls-1],&NNf); CHKERRQ(ierr);
	ierr = VecCopy(NN,NNf); CHKERRQ(ierr);
	
	// VecTagger is used to find indices i where NN[i] = 1
	ierr = VecTaggerCreate(PETSC_COMM_WORLD, &tagger); CHKERRQ(ierr);
	ierr = VecTaggerSetType(tagger, VECTAGGERABSOLUTE); CHKERRQ(ierr);
	taggerbox.min = 0.9; taggerbox.max = 1.1;
	ierr = VecTaggerAbsoluteSetBox(tagger, &taggerbox); CHKERRQ(ierr);
	ierr = VecTaggerSetUp(tagger); CHKERRQ(ierr);
}
// END OF NEW CODE		

for (PetscInt k = 1; k < nlvls; k++) 
{			            
	if (pcmgredistribute==PETSC_FALSE) {
		// ORIGINAL CODE
		ierr = DMCreateInterpolation(da_list[k - 1], da_list[k], &R, NULL); CHKERRQ(ierr);
		ierr = PCMGSetInterpolation(ipc, k, R); CHKERRQ(ierr);			
		ierr = MatDestroy(&R); CHKERRQ(ierr);
		// END OF ORIGINAL CODE
	}				
	else if (pcmgredistribute==PETSC_TRUE) 
	{
		// NEW CODE 
		//
		// Partially inspired by https://petsc.org/release/src/snes/impls/vi/rs/virs.c.html 
		//
		// NOTES: * Enable from command line with "-S_pc_type redistribute -pcmgredistribute 1"   
		//          (without ""). The OLD CODE runs by default
		//        * Right now we have hardcoded for just one level
		//        * So far this works only on one rank. To make it work on multiple ranks we must shuffle
		//          around the rows and and cols following the order of redistribute.c
		// 
		
		// Variables for convenience (to avoid confusing which is coarse and which is fine)
		DM *dmc = &(da_list[k-1]);	// Coarse mesh
		DM *dmf = &(da_list[k]);    // Fine mesh
		
		ierr = DMCreateInterpolation(*dmc, *dmf, &R, NULL); CHKERRQ(ierr);
		
		PetscInt m,n;
		ierr = MatGetSize(R, &m, &n); CHKERRQ(ierr);
		PetscPrintf(PETSC_COMM_WORLD,"%s, %d: \nk=%d, size(R)=(%d,%d)\n",__FILE__,__LINE__,k,m,n);
		
		// Create IS with freedofs on the finer grid
		ierr = VecTaggerComputeIS(tagger, NNf, &isfreef, NULL); CHKERRQ(ierr);
		
		// Create IS with freedofs on the coarse grid. This is a subset of the freedofs on the fine grid
		// obtained by simply selecting the values on the coarse grid points.
		Mat inject;
		ierr = DMCreateGlobalVector(*dmc,&NNc); CHKERRQ(ierr);
		ierr = DMCreateInjection(*dmc, *dmf, &inject); CHKERRQ(ierr);  // Make a version of NN on the coarse grid
		ierr = MatRestrict(inject, NNf, NNc);  CHKERRQ(ierr);				
		ierr = MatDestroy(&inject); CHKERRQ(ierr);				
		ierr = VecTaggerComputeIS(tagger, NNc, &isfreec, NULL); CHKERRQ(ierr);
		
		ierr = VecGetSize(NNc,&m); CHKERRQ(ierr);
		PetscPrintf(PETSC_COMM_WORLD,"%s, %d: \nsize(NNc)=(%d,1)\n",__FILE__,__LINE__,m);
								
		// Keep only rows and cols that are free on each level
		ierr = MatCreateSubMatrix(R, isfreef, isfreec, MAT_INITIAL_MATRIX, &Rred); CHKERRQ(ierr);
		ierr = PCMGSetInterpolation(ipc, k, Rred); CHKERRQ(ierr);			
		
		// Print to Matlab and clean up
		ierr = PrintMatToFile(K,"K"); CHKERRQ(ierr);
		ierr = PrintISToFile(isfreef,"isfreef"); CHKERRQ(ierr);
		ierr = PrintISToFile(isfreec,"isfreec"); CHKERRQ(ierr);
		ierr = PrintMatToFile(R,"R"); CHKERRQ(ierr);
		ierr = PrintMatToFile(Rred,"Rred"); CHKERRQ(ierr);
		
		ierr = ISDestroy(&isfreef); CHKERRQ(ierr);
		ierr = ISDestroy(&isfreec); CHKERRQ(ierr);
		ierr = MatDestroy(&Rred); CHKERRQ(ierr);
		ierr = MatDestroy(&R); CHKERRQ(ierr);
		ierr = VecDestroy(&NNf); CHKERRQ(ierr);
		ierr = VecDestroy(&NNc); CHKERRQ(ierr);
		ierr = VecTaggerDestroy(&tagger); CHKERRQ(ierr);
		// END OF NEW CODE
	}		
}