/*!\file:  MeshPartition.cpp
 * \brief partition elements and grids across a cluster of size numprocs. 
 */


#include "./MeshPartitionx.h"

#undef __FUNCT__ 
#define __FUNCT__ "MeshPartitionx"

#include "../shared/shared.h"
#include "../include/macros.h"
#include "../toolkits/toolkits.h"
#include "../EnumDefinitions/EnumDefinitions.h"


int MeshPartitionx(int** pepart, int** pnpart, int numberofelements,int numberofgrids,double* elements,
		int numberofelements2d,int numberofgrids2d,double* elements2d,int numlayers,int elements_width, char* meshtype,int num_procs){

	int noerr=1;
	int i,j;

	/*Metis partitioning: */
	int* epart=NULL;
	int* npart=NULL;
	int* index=NULL;

	int* epart2d=NULL;
	int* npart2d=NULL;
	int* index2d=NULL;
	int  count=0;

	int  etype=1; //tria mesh see metis/Programs/Io.c
	int  etype2d=1; //tria mesh see metis/Programs/Io.c
	int  numflag=0;
	int  edgecut=1;

	if(strcmp(meshtype,"2d")==0){
		epart=(int*)xmalloc(numberofelements*sizeof(int));
		npart=(int*)xmalloc(numberofgrids*sizeof(int));
		index=(int*)xmalloc(elements_width*numberofelements*sizeof(int));
		for (i=0;i<numberofelements;i++){
			for (j=0;j<elements_width;j++){
				*(index+elements_width*i+j)=(int)*(elements+elements_width*i+j)-1; //-1 for C indexing in Metis
			}
		}

		/*Partition using Metis:*/
		METIS_PartMeshNodal(&numberofelements,&numberofgrids, index, &etype, &numflag, &num_procs, &edgecut, epart, npart);
	}
	else{
		/*We have a 3d mesh, made of a regularly extruded 2d mesh. We first partition the 2d mesh, then we extrude the partition: */

		/*First build concatenated 2d mesh  from 2d_coll and 2d_noncoll: */
		epart2d=(int*)xmalloc(numberofelements2d*sizeof(int));
		npart2d=(int*)xmalloc(numberofgrids2d*sizeof(int)); 
		index2d=(int*)xmalloc(3*numberofelements2d*sizeof(int));

		for (i=0;i<numberofelements2d;i++){
			for (j=0;j<3;j++){
				*(index2d+3*i+j)=(int)*(elements2d+3*i+j)-1; //-1 for C indexing in Metis
			}
		}

		METIS_PartMeshNodal(&numberofelements2d,&numberofgrids2d, index2d, &etype2d, &numflag, &num_procs, &edgecut, epart2d, npart2d);

		/*Extrude epart2d to epart, using numlayers: */
		epart=(int*)xmalloc(numberofelements*sizeof(int));
		
		count=0;
		for(i=0;i<(numlayers-1);i++){
			for(j=0;j<numberofelements2d;j++){
				epart[count]=epart2d[j];
				count++;
			}
		}
		throw ErrorException(__FUNCT__," grid partitioning not implemented yet in 3d");

	}

	#ifdef _DEBUG_
	if(my_rank==0){
		for (i=0;i<numberofelements;i++){
			printf("El %i My rank %i\n",i+1,epart[i]);
		}
	}
	#endif
	
	/*Assign output pointer:*/
	*pepart=epart;
	*pnpart=npart;

	/*Free ressources: */
	xfree((void**)&index);
	xfree((void**)&epart2d);
	xfree((void**)&npart2d);
	xfree((void**)&index2d);

	return noerr;

}
