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

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

int MeshPartitionx(int** pepart, int** pnpart, int numberofelements,int numberofnodes,IssmDouble* elements,
		int numberofelements2d,int numberofnodes2d,IssmDouble* elements2d,int numlayers,int elements_width, int dim,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(dim==2){
		epart=xNew<int>(numberofelements);
		npart=xNew<int>(numberofnodes);
		index=xNew<int>(elements_width*numberofelements);
		for (i=0;i<numberofelements;i++){
			for (j=0;j<elements_width;j++){
				*(index+elements_width*i+j)=reCast<int>(*(elements+elements_width*i+j))-1; //-1 for C indexing in Metis
			}
		}

		/*Partition using Metis:*/
		if (num_procs>1){
			#ifdef _HAVE_METIS_
			METIS_PartMeshNodalPatch(&numberofelements,&numberofnodes, index, &etype, &numflag, &num_procs, &edgecut, epart, npart);
			#endif
		}
		else if (num_procs==1){
			/*METIS does not know how to deal with one cpu only!*/
			for (i=0;i<numberofelements;i++) epart[i]=0;
			for (i=0;i<numberofnodes;i++)    npart[i]=0;
		}
		else _error2_("At least one processor is required");
	}
	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=xNew<int>(numberofelements2d);
		npart2d=xNew<int>(numberofnodes2d); 
		index2d=xNew<int>(3*numberofelements2d);

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

		/*Partition using Metis:*/
		if (num_procs>1){
			#ifdef _HAVE_METIS_
			METIS_PartMeshNodalPatch(&numberofelements2d,&numberofnodes2d, index2d, &etype2d, &numflag, &num_procs, &edgecut, epart2d, npart2d);
			#endif
		}
		else if (num_procs==1){
			/*METIS does not know how to deal with one cpu only!*/
			for (i=0;i<numberofelements2d;i++) epart2d[i]=0;
			for (i=0;i<numberofnodes2d;i++)    npart2d[i]=0;
		}
		else _error2_("At least one processor is required");

		/*Extrude epart2d to epart, using numlayers: */
		epart=xNew<int>(numberofelements);
		
		count=0;
		for(i=0;i<(numlayers-1);i++){
			for(j=0;j<numberofelements2d;j++){
				epart[count]=epart2d[j];
				count++;
			}
		}

		/*Extrude npart2d to npart, using numlayers: */
		npart=xNew<int>(numberofnodes);
		
		count=0;
		for(i=0;i<(numlayers);i++){
			for(j=0;j<numberofnodes2d;j++){
				npart[count]=npart2d[j];
				count++;
			}
		}
	}
	
	/*Assign output pointer:*/
	*pepart=epart;
	*pnpart=npart;

	/*Free ressources: */
	xDelete<int>(index);
	xDelete<int>(epart2d);
	xDelete<int>(npart2d);
	xDelete<int>(index2d);
	return noerr;
}
