//
// Software License for AMDiS
//
// Copyright (c) 2010 Dresden University of Technology 
// All rights reserved.
// Authors: Simon Vey, Thomas Witkowski et al.
//
// This file is part of AMDiS
//
// See also license.opensource.txt in the distribution.

#include <boost/lexical_cast.hpp>
#include "parallel/ZoltanPartitioner.h"
#include "Traverse.h"
#include "ElInfo.h"

namespace AMDiS {

  ZoltanPartitioner::ZoltanPartitioner(MPI::Intracomm *comm)
    : MeshPartitioner(comm),
      zoltan(*comm)
  {
    zoltan.Set_Num_Obj_Fn(ZoltanFunctions::getNumObj, this);
    zoltan.Set_Obj_List_Fn(ZoltanFunctions::getObjectList, this);
    zoltan.Set_Num_Geom_Fn(ZoltanFunctions::getNumGeom, this);
    zoltan.Set_Geom_Multi_Fn(ZoltanFunctions::getGeomMulti, this);
  }


  bool ZoltanPartitioner::partition(map<int, double> &elemWeights,
				    PartitionMode mode)
  {
    FUNCNAME("ZoltanPartitioner::partition()");

    using boost::lexical_cast;

    int changes;
    int nGid, nLid;

    int nImportEls;   
    ZOLTAN_ID_PTR import_global_ids, import_local_ids;
    int *import_procs;
    int *import_to_part;

    int nExportEls;
    ZOLTAN_ID_PTR export_global_ids, export_local_ids;
    int *export_procs;
    int *export_to_part;

    zoltan.Set_Param("LB_METHOD", "RCB");
    zoltan.Set_Param("NUM_GLOBAL_PARTS", 
		     lexical_cast<string>(mpiComm->Get_size()).c_str());
    int err = zoltan.LB_Partition(changes, nGid, nLid, 
				  nImportEls,
				  import_global_ids,
				  import_local_ids,
				  import_procs,
				  import_to_part,
				  nExportEls,
				  export_global_ids,
				  export_local_ids,
				  export_procs,
				  export_to_part);
    
    recvElements.clear();
    sendElements.clear();

    if (err == ZOLTAN_OK && changes != 0) {   
      if (nImportEls > 0) {
	for (int i = 0; i < nImportEls; i++) {
	  int recvElIndex = import_global_ids[i];
	  
	  elementInRank[recvElIndex] = true;
	  recvElements[import_procs[i]].push_back(recvElIndex);
	}
      }
      
      if (nExportEls > 0) {
	for (int i = 0; i < nExportEls; i++) {
	  int sendElIndex = export_global_ids[i];
	  
	  elementInRank[sendElIndex] = false;
	  sendElements[export_procs[i]].push_back(sendElIndex);
	}      
      }
    }

    zoltan.LB_Free_Part(&import_global_ids, &import_local_ids, 
			&import_procs, &import_to_part);
    zoltan.LB_Free_Part(&export_global_ids, &export_local_ids, 
			&export_procs, &export_to_part);      

    if (err != ZOLTAN_OK)
      return false;

    return true;
  }


  void ZoltanPartitioner::getPartitionMap(map<int, int> &partitionMap)
  {
    FUNCNAME("ZoltanPartitioner::getPartitionMap()");

    int mpiSize = mpiComm->Get_size();

    vector<int> localElements;
    for (map<int, bool>::iterator it = elementInRank.begin(); 
	 it != elementInRank.end(); ++it)
      if (it->second)
	localElements.push_back(it->first);
 
    int nLocalElements = localElements.size();
    vector<int> nPartitionElements(mpiSize);
    vector<int> elDist(mpiSize + 1);
    mpiComm->Allgather(&nLocalElements, 1, MPI_INT, &(elDist[1]), 1, MPI_INT);

    elDist[0] = 0;
    nPartitionElements[0] = elDist[1];

    for (int i = 2; i <= mpiSize; i++) {
      nPartitionElements[i - 1] = elDist[i];
      elDist[i] += elDist[i - 1];
    }

    int nOverallElements = elDist[mpiSize];

    TEST_EXIT_DBG(nOverallElements == static_cast<int>(elementInRank.size()))
      ("Number of elements differs: %d %d!\n", nOverallElements, elementInRank.size());

    vector<int> partitionElements(nOverallElements);

    // distribute partition elements
    mpiComm->Allgatherv(&(localElements[0]),
			nLocalElements, 
			MPI_INT, 
			&(partitionElements[0]), 
			&(nPartitionElements[0]), 
			&(elDist[0]),
			MPI_INT);

    // fill partitionMap
    for (int i = 0; i < mpiSize; i++)
      for (int j = 0; j < nPartitionElements[i]; j++)
	partitionMap[partitionElements[elDist[i] + j]] = i;
  }

  

  int ZoltanFunctions::getNumObj(void *data, int *ierr)
  {
    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
    map<int, bool>& elInRank = zoltan->getElementInRank();

    int nObjects = 0;
    for (map<int, bool>::iterator it = elInRank.begin(); it != elInRank.end(); ++it)
      if (it->second)
	nObjects++;

    *ierr = ZOLTAN_OK;

    return nObjects;
  }


  void ZoltanFunctions::getObjectList(void *data, 
				      int sizeGid, 
				      int sizeLid,
				      ZOLTAN_ID_PTR globalId, 
				      ZOLTAN_ID_PTR localId,
				      int wgt_dim, 
				      float *obj_wgts, 
				      int *ierr)
  {
    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
    map<int, bool>& elInRank = zoltan->getElementInRank();
    int localCounter = 0;

    for (map<int, bool>::iterator it = elInRank.begin(); it != elInRank.end(); ++it) {
      if (it->second) {
	globalId[localCounter] = it->first;
	localId[localCounter] = localCounter;
	localCounter++;
      }
    }

    *ierr = ZOLTAN_OK;
  }


  int ZoltanFunctions::getNumGeom(void *data, int *ierr)
  {
    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
    *ierr = ZOLTAN_OK;

    return zoltan->getMesh()->getGeo(WORLD);
  }


  void ZoltanFunctions::getGeomMulti(void *data, 
				     int num_gid_entries, 
				     int num_lid_entries, 
				     int num_obj, 
				     ZOLTAN_ID_PTR global_ids, 
				     ZOLTAN_ID_PTR local_ids, 
				     int num_dim, 
				     double *geom_vec, 
				     int *ierr)
  {
    FUNCNAME("ZoltanFunctions::getGeomMulti()");

    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
    Mesh *mesh = zoltan->getMesh();

    TEST_EXIT_DBG(num_dim == mesh->getGeo(WORLD))("Should not happen!\n");

    map<int, WorldVector<double> > elIndexToCoords;    
    DimVec<double> bary(mesh->getDim(), DEFAULT_VALUE, 1.0 / mesh->getGeo(VERTEX));
    WorldVector<double> coords;
    TraverseStack stack;
    ElInfo *elInfo = 
      stack.traverseFirst(mesh, 0, Mesh::CALL_EL_LEVEL | Mesh::FILL_COORDS);
    while (elInfo) {
      elInfo->coordToWorld(bary, coords);
      elIndexToCoords[elInfo->getElement()->getIndex()] = coords;

      elInfo = stack.traverseNext(elInfo);
    }

    int c = 0;
    for (int i = 0; i < num_obj; i++) {
      int elIndex = global_ids[i];

      TEST_EXIT_DBG(elIndexToCoords.count(elIndex))("Should not happen!\n");
      
      for (int j = 0; j < num_dim; j++)
	geom_vec[c++] = elIndexToCoords[elIndex][j];
    }

    *ierr = ZOLTAN_OK;
  }
  
}