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

#undef __FUNCT__ 
#define __FUNCT__ "PartitionSets"

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

	int i;

	/*output: */
	Vec partitionb=NULL;
	double* partitionb_local=NULL;
	Vec partitionc=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]) throw ErrorException(__FUNCT__,exprintf("%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);
	VecCreateMPIWithArray(MPI_COMM_WORLD,csize,PETSC_DECIDE,partitionc_local,&partitionc);

	/*Free ressources:*/
	xfree((void**)&idxm);
	xfree((void**)&flags_a_local);
	xfree((void**)&flags_b_local);
	xfree((void**)&flags_c_local);

	VecAssemblyBegin(partitionb); 
	VecAssemblyEnd(partitionb);

	VecAssemblyBegin(partitionc); 
	VecAssemblyEnd(partitionc);

	/*Assign output pointers*/
	*ppartitionb=partitionb;
	*ppartitionc=partitionc;


}
