#include "MeshPartition_METIS5.h"
#include <iostream>

MeshPartition_METIS5::MeshPartition_METIS5(){
   m_numCuts = 0;
   m_numNodes = 0;
   m_numTets = 0;

   m_tetList = 0;
   m_nodeList = 0;

   m_elemNodeIDs = 0;
   m_elemNodeWeights = 0;
   m_elemPtr = 0;
   m_elemPtrCutted = 0;
   m_elemNodeIDsCutted = 0;
}

MeshPartition_METIS5::~MeshPartition_METIS5(){
   if(m_elemNodeIDs){ 
      delete [] m_elemNodeIDs;
      m_elemNodeIDs = 0;
   }
   if(m_elemNodeWeights){ 
      delete [] m_elemNodeWeights;
      m_elemNodeWeights = 0;
   }
   if(m_elemPtr){ 
      delete [] m_elemPtr;
      m_elemPtr = 0;
   }
   if(m_elemPtrCutted){ 
      delete [] m_elemPtrCutted;
      m_elemPtrCutted = 0;
   }
   if(m_elemNodeIDsCutted){ 
      delete [] m_elemNodeIDsCutted; 
      m_elemNodeIDsCutted = 0;
   }
}

void MeshPartition_METIS5::PerformSetup(int numNodes, int numTets, node * nodeList, tetra * tetList, int numCuts){
   m_numNodes = numNodes;
   m_numTets  = numTets;
   m_nodeList = nodeList;
   m_tetList  = tetList;
   m_numCuts  = numCuts;
}

////////////////////////////////////////////////////////////////////////////////////////////
/// This function loop through the partition list, according to the number of cuts stored for
/// each partition, perform the cut through METIS
/// \note
///      - 
/// \version
///      - Zhao Bo, 30/Jan/2012, created
////////////////////////////////////////////////////////////////////////////////////////////

void MeshPartition_METIS5::PerformCut (std::vector < fp_t > & tetWeightList){
   if(m_numCuts > 1){
      int objval(0);

      m_elemNodeIDs = new int[4 * m_numTets];
      m_elemNodeWeights = new int[m_numNodes];
      m_elemPtr = new int[m_numTets + 1];
      m_elemPtrCutted = new int[m_numTets];
      m_elemNodeIDsCutted = new int[4 * m_numTets];

      memset (m_elemNodeWeights, 0, m_numNodes * sizeof (int));

      for(int j = 0; j < m_numTets; j++){
         tetra & tetj = m_tetList[j];
         for(int k = 0; k < 4; k++){
            int nodeID = tetj.getNode(k)->getid();
            m_elemNodeIDs[j * 4 + k] = nodeID;
            m_elemNodeWeights[nodeID] += (int) tetWeightList[j];
            // m_elemNodeWeights[nodeID] = 1;
         }
         m_elemPtr[j] = j * 4;
      }
      m_elemPtr[m_numTets] = m_numTets * 4;

      int info = METIS_PartMeshNodal(&m_numTets, &m_numNodes, m_elemPtr, m_elemNodeIDs, m_elemNodeWeights, NULL, &m_numCuts, NULL, NULL, &objval, m_elemPtrCutted, m_elemNodeIDsCutted);
      if(info != 1)
         std::cout << "MeshPartition_METIS5" << __FILE__ << __LINE__ << "METIS5 function failed !!!\n";
   }
}

void MeshPartition_METIS5::PerformCut(){
   if(m_numCuts > 1){
      int objval(0);

      m_elemNodeIDs = new int[4 * m_numTets];
      m_elemPtr = new int[m_numTets + 1];
      m_elemPtrCutted = new int[m_numTets];
      m_elemNodeIDsCutted = new int[4 * m_numTets];

      for(int j = 0; j < m_numTets; j++){
         tetra & tetj = m_tetList[j];
         for(int k = 0; k < 4; k++){
            int nodeID = tetj.getNode(k)->getid();
            m_elemNodeIDs[j * 4 + k] = nodeID;
         }
         m_elemPtr[j] = j * 4;
      }
      m_elemPtr[m_numTets] = m_numTets * 4;
//       cout << "here 2" << endl;
      int info = METIS_PartMeshNodal(&m_numTets, &m_numNodes, m_elemPtr , m_elemNodeIDs, NULL, NULL, &m_numCuts , NULL, NULL, &objval, m_elemPtrCutted, m_elemNodeIDsCutted);

      if(info != 1)
         std::cout << "MeshPartition_METIS5" << __FILE__ << __LINE__ << "METIS5 function failed !!!\n";
   }
}


void MeshPartition_METIS5::PerformSetup(int numNodes, int numTets, int numCuts){
   m_numNodes = numNodes;
   m_numTets = numTets;
   m_numCuts = numCuts;
}

void MeshPartition_METIS5::PerformCut(int* elmnts){
   if(m_numCuts > 1){
      int objval(0);

      m_elemNodeIDs = new int[4 * m_numTets];
      m_elemPtr = new int[m_numTets + 1];
      m_elemPtrCutted = new int[m_numTets];
      m_elemNodeIDsCutted = new int[4 * m_numTets];

      for(int j = 0; j < m_numTets; j++) {
         for(int k = 0; k < 4; k++){
            m_elemNodeIDs[j*4+k] = elmnts[j*4+k];
         }
         m_elemPtr[j] = j * 4;
      }
      m_elemPtr[m_numTets] = m_numTets * 4;

      int info = METIS_PartMeshNodal (&m_numTets, &m_numNodes, m_elemPtr , m_elemNodeIDs, NULL, NULL, &m_numCuts , NULL, NULL, &objval, m_elemPtrCutted, m_elemNodeIDsCutted);

      if(info != 1)
         std::cout << "MeshPartition_METIS5" << __FILE__ << __LINE__ << "METIS5 function failed !!!\n";
   }
}

int* MeshPartition_METIS5::PartionGraphForGrouping(vector<int>& vec_xadj, vector<int>& vec_adjncy){
   int objval  = 0;

	idx_t option[METIS_NOPTIONS];
	METIS_SetDefaultOptions(option);
	option[METIS_OPTION_OBJTYPE] = METIS_OBJTYPE_CUT;
	option[METIS_OPTION_PTYPE] = METIS_PTYPE_KWAY;
	option[METIS_OPTION_CTYPE] = METIS_CTYPE_RM; 
	option[METIS_OPTION_NCUTS] = 1; 
	option[METIS_OPTION_UFACTOR] = 30; 
	option[METIS_OPTION_NUMBERING] = 0; 
	option[METIS_OPTION_CONTIG] = 1;
	option[METIS_OPTION_DBGLVL] = 0;

   int ncon = 1;
   int VertNum = ((int)vec_xadj.size()-1);
   int* Partion = new int[VertNum];
   int* xadj = new int[(int)vec_xadj.size()];
   int* adjncy = new int[(int)vec_adjncy.size()];

   for(int idx = 0; idx < (int)vec_xadj.size()   ; idx++){ 
       xadj[idx] = vec_xadj[idx]; 
   }
   for(int idx = 0; idx < (int)vec_adjncy.size() ; idx++){
       adjncy[idx] = vec_adjncy[idx]; 
   }

   int metis_retun =  METIS_PartGraphKway(&VertNum, &ncon, xadj, adjncy,NULL, NULL, NULL, &m_numCuts, NULL, NULL, option, &objval, Partion);
 
   if(metis_retun !=  METIS_OK){
      cout << "FemGrpBalancing::PerformGroupingMetis->METIS_PartGraphKway" << __FILE__ << __LINE__ << "METIS5 function failed !!!\n";
   }

   delete [] xadj;
   delete [] adjncy;
   return Partion;
}