/*!\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"

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

		
	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;

		if (range){
			node_rows=xNew<int>(range); //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=xNew<double>(count); //holder for the values to be extracted from vectorA
		}
		else{
			xDelete<int>(node_rows); //count=0 means no values was condensed out for this node. null node_rows for use in VecGetValues.
			values=NULL;
		}

		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);

		/*build new node_rows index.*/
		if (count){
			for (i=0;i<count;i++){
				*(node_rows+i)=lower_row+i;
			}
		}
		
		/* outvector should not be partitioned like it was previously, but according to row_partition_vector_size, this in case we 
		 * are running with the special kffpartition schema: */
		if(kffpartition){
			VecFree(&outvector);
			outvector=NewVec(row_partition_vector_size);
		}

		if (count){
			VecSetValues(outvector,count,node_rows,values,INSERT_VALUES);
		}
		
		/*Assemble vector*/
		VecAssemblyBegin(outvector);
		VecAssemblyEnd(outvector);
	}
	
	/*Assign output pointers:*/
	*poutvector=outvector;
	xDelete<int>(node_rows);
	xDelete<double>(values);
}
