/*!\file:  VecPartition.cpp
 * \brief  partition vector according to a node set
 */ 

#ifdef HAVE_CONFIG_H
	#include "config.h"
#else
#error "Cannot compile with HAVE_CONFIG_H symbol! run configure first!"
#endif

/*Petsc includes: */
#include "petscmat.h"
#include "petscvec.h"
#include "petscksp.h"

#include "../../../shared/shared.h"


#undef __FUNCT__ 
#define __FUNCT__ "VecPartition"

void VecPartition(Vec* poutvector,Vec vectorA, double* row_partition_vector, int row_partition_vector_size){

		
	int i;
	
	/*Petsc matrix*/
	int* node_rows=NULL;
	double* values=NULL;
	int count;
	int lower_row,upper_row,range;

	/*Input*/
	int MA; //Vector dimension

	/*output*/
	Vec outvector=NULL;

	
	/*Get size of input vector: */
	VecGetSize(vectorA,&MA);
	
    /*Condense out*/

    /*If the partitioning row vector has a 0 dimension, we return a NILL vector*/
    if (row_partition_vector_size==0){
		VecCreate(PETSC_COMM_WORLD,&outvector);   
		VecSetSizes(outvector,PETSC_DECIDE,0);
		VecSetFromOptions(outvector);
		
		VecAssemblyBegin(outvector);
		VecAssemblyEnd(outvector);
	}
	else{
		/*Locally get values and build a new vector which is a subvector*/
		
		VecGetOwnershipRange(vectorA,&lower_row,&upper_row);
		upper_row--;
		range=upper_row-lower_row+1;
		#ifdef _DEBUG_
		 	PetscSynchronizedPrintf(MPI_COMM_WORLD,"My rank %i Range %i\n",my_rank,range);
			PetscSynchronizedFlush(MPI_COMM_WORLD);
		#endif

		if (range){
			node_rows=(int*)xmalloc(range*sizeof(int)); //this is the maximum number of rows one node can extract.
		
			count=0;
			for (i=0;i<row_partition_vector_size;i++){
				if ( ((int)(*(row_partition_vector+i))>=(lower_row+1)) && ((int)(*(row_partition_vector+i))<=(upper_row+1)) ){
					*(node_rows+count)=(int)*(row_partition_vector+i)-1;
					count++;
				}
			}
		}
		else{
			count=0;
		}
					
     	if (count){
			values=(double*)xmalloc(count*sizeof(double)); //holder for the values to be extracted from vectorA
		}
		else{
			xfree((void**)&node_rows); //count=0 means no values was condensed out for this node. null node_rows for use in VecGetValues.
			values=NULL;
		}

		#ifdef _DEBUG_
			PetscSynchronizedPrintf(MPI_COMM_WORLD,"My rank: %i My count: %i node_rows: %p values: %p\n",my_rank,count,node_rows,values);
			PetscSynchronizedFlush(MPI_COMM_WORLD);
		#endif 

		if (count){
			VecGetValues(vectorA,count,node_rows,values);
		}
		
		/*Ok, each node now have count values corresponding to the node_rows extracted values from vectorA.
		 * From count and values, create a new vector*/

		VecCreate(PETSC_COMM_WORLD,&outvector);    ;
		VecSetSizes(outvector,count,PETSC_DECIDE);
		VecSetFromOptions(outvector);

		VecGetOwnershipRange(outvector,&lower_row,&upper_row);
		upper_row--;
		range=upper_row-lower_row+1;

		/*Set values using the values vector. First build new node_rows index.*/
		if (count){
			for (i=0;i<count;i++){
				*(node_rows+i)=lower_row+i;
			}
		}
		if (count){
			VecSetValues(outvector,count,node_rows,values,INSERT_VALUES);
		}
		
		/*Assemble vector*/
		VecAssemblyBegin(outvector);
		VecAssemblyEnd(outvector);

		
	}
	
	#ifdef _DEBUG_
		_printf_("Vector A:\n");
		fflush(stdout);
		VecView(vectorA,PETSC_VIEWER_STDOUT_WORLD);
		_printf_("Partition of vector A:\n");
		fflush(stdout);
		VecView(outvector,PETSC_VIEWER_STDOUT_WORLD);
	#endif
	
	/*Assign output pointers:*/
	*poutvector=outvector;

	/*Free ressources:*/
	xfree((void**)&node_rows);
	xfree((void**)&values);

}
