/*!\file:  PartitionSets.cpp
 * \brief partition degree of freedome sets
 */ 
	
#include "stdio.h"
#include "../../shared/shared.h"
#include "../../include/include.h"

void PartitionSets(Vec* ppartitionb, Vec* ppartitionc,Vec flags_a,Vec flags_b,Vec flags_c,int gsize){

	int i;

	/*output: */
	Vec partitionb=NULL;
	Vec partitionb_vec=NULL;
	double* partitionb_local=NULL;
	Vec partitionc=NULL;
	Vec partitionc_vec=NULL;
	double* partitionc_local=NULL;

	/*intermediary: */
	int  local_size;
	int* idxm=NULL;
	double* flags_a_local=NULL;
	double* flags_b_local=NULL;
	double* flags_c_local=NULL;
	int     lower_row, upper_row, range;
	int     asize, bsize, csize;
	int     acount, bcount, ccount, dummy;

	

	/*First recover flags locally on each cpu: */
	VecGetOwnershipRange(flags_a,&lower_row,&upper_row); upper_row--; range=upper_row-lower_row+1;    

	if(range){
		idxm=(int*)xmalloc(range*sizeof(int));
		flags_a_local=(double*)xmalloc(range*sizeof(double));
		flags_b_local=(double*)xmalloc(range*sizeof(double));
		flags_c_local=(double*)xmalloc(range*sizeof(double));
	}
	for(i=0;i<range;i++)idxm[i]=lower_row+i;

	VecGetValues(flags_a,range,idxm,flags_a_local);
	VecGetValues(flags_b,range,idxm,flags_b_local);
	VecGetValues(flags_c,range,idxm,flags_c_local);


	/*Figure out local size of a, b and c sets on each cpu: */
	asize=0;
	bsize=0;
	csize=0;
	for(i=0;i<range;i++){
		if (flags_a_local[i]){
			if (flags_b_local[i])bsize++;
			if (flags_c_local[i])csize++;
			asize++;
		}
	}

	/*Figure out start of "a" set local cpu counter: */
	GetOwnershipBoundariesFromRange(&acount,&dummy,asize);

	/*Allocate local partitioning vectors: */
	if(bsize)partitionb_local=(double*)xmalloc(bsize*sizeof(double));
	if(csize)partitionc_local=(double*)xmalloc(csize*sizeof(double));

	/*Loop again, and plug values: */
	bcount=0;
	ccount=0;
	for(i=0;i<range;i++){
		if (flags_a_local[i]){
			if (flags_b_local[i]){
				partitionb_local[bcount]=acount+1; //matlab indexing
				bcount++;
			}
			if (flags_c_local[i]){
				partitionc_local[ccount]=acount+1; //matlab indexing
				ccount++;
			}
			if (flags_b_local[i] && flags_c_local[i]) _error_("%s%i%s"," for dof ",i,": breach of exclusive partitioning between sets");
			
			acount++;
		}
	}

	/*Now, using the local partitions for b and c, create parallel vectors: */
	VecCreateMPIWithArray(MPI_COMM_WORLD,bsize,PETSC_DECIDE,partitionb_local,&partitionb_vec);
	VecCreateMPIWithArray(MPI_COMM_WORLD,csize,PETSC_DECIDE,partitionc_local,&partitionc_vec);

	VecAssemblyBegin(partitionb_vec); 
	VecAssemblyEnd(partitionb_vec);

	VecAssemblyBegin(partitionc_vec); 
	VecAssemblyEnd(partitionc_vec);

	/*Now we must duplicate these vectors because:
	 *
	 * Petsc specifies on http://www.mcs.anl.gov/petsc/petsc-as/snapshots/petsc-current/docs/manualpages/Vec/VecCreateMPIWithArray.html
	 * "PETSc does NOT free the array when the vector is destroyed via VecDestroy().
	 * The user should not free the array until the vector is destroyed"
	 *
	 * The only way we can keep the vector partitionb_vec while destroying the array partitionb_local
	 * is to duplicate the vector and then destroy the initial vector as well as the array*/
	VecDuplicate(partitionb_vec,&partitionb); VecCopy(partitionb_vec,partitionb); VecFree(&partitionb_vec); xfree((void**)&partitionb_local);
	VecDuplicate(partitionc_vec,&partitionc); VecCopy(partitionc_vec,partitionc); VecFree(&partitionc_vec); xfree((void**)&partitionc_local);

	/*Free ressources:*/
	xfree((void**)&idxm);
	xfree((void**)&flags_a_local);
	xfree((void**)&flags_b_local);
	xfree((void**)&flags_c_local);
	
	/*Assign output pointers*/
	*ppartitionb=partitionb;
	*ppartitionc=partitionc;

}