#include "ParallelDomainProblem.h"
#include "ProblemScal.h"
#include "ProblemInstat.h"
#include "ParMetisPartitioner.h"
#include "Mesh.h"
#include "Traverse.h"
#include "ElInfo.h"
#include "Element.h"
#include "MacroElement.h"
#include "PartitionElementData.h"

namespace AMDiS {

  ParallelDomainProblemBase::ParallelDomainProblemBase(const std::string& name,
						       ProblemIterationInterface *iIF,
						       ProblemTimeInterface *tIF,
						       FiniteElemSpace *fe)
    : iterationIF(iIF),
      timeIF(tIF),
      feSpace(fe),
      mesh(fe->getMesh()),
      initialPartitionMesh(true),
      nRankDOFs(0)
  {
    mpiRank = MPI::COMM_WORLD.Get_rank();
    mpiSize = MPI::COMM_WORLD.Get_size();
    mpiComm = MPI::COMM_WORLD;
    partitioner = new ParMetisPartitioner(mesh, &mpiComm);
  }

  Flag ParallelDomainProblemBase::oneIteration(AdaptInfo *adaptInfo, Flag toDo)
  {
    return iterationIF->oneIteration(adaptInfo, toDo);
  }

  void ParallelDomainProblemBase::initParallelization(AdaptInfo *adaptInfo)
  {
    if (mpiSize <= 1)
      return;

    // create an initial partitioning of the mesh
    partitioner->createPartitionData();
    // set the element weights, which are 1 at the very first begin
    setElemWeights(adaptInfo);
    // and now partition the mesh
    partitionMesh(adaptInfo);   

    std::map<int, std::set<int> > partitionDofs;

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, Mesh::CALL_LEAF_EL);
    int nLeaves = 0;
    while (elInfo) {
      Element *element = elInfo->getElement();

      // Hidde elements which are not part of ranks partition.
      PartitionElementData *partitionData = 
	dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));
      if (partitionData->getPartitionStatus() != IN) {
      } else {
      }

      // Determine to each dof the partition(s) it corresponds to.
      for (int i = 0; i < 3; i++) 
	partitionDofs[element->getDOF(i, 0)].insert(partitionVec[element->getIndex()]);
          
      elInfo = stack.traverseNext(elInfo);
    }

    std::vector<int> rankDofs;
    for (std::map<int, std::set<int> >::iterator it = partitionDofs.begin();
	 it != partitionDofs.end();
	 ++it) {
      for (std::set<int>::iterator itpart1 = it->second.begin();
	   itpart1 != it->second.end();
	   ++itpart1) {
	if (*itpart1 == mpiRank) {
	  if (it->second.size() == 1) {
	    rankDofs.push_back(it->first);
	  } else {	    
	    bool insert = true;
	    for (std::set<int>::iterator itpart2 = it->second.begin();
		 itpart2 != it->second.end();
		 ++itpart2) {
	      if (*itpart2 > mpiRank) {
		insert = false;
		break;
	      }
	    }
	    if (insert) {
	      rankDofs.push_back(it->first);
	    }
	  }
	}
      }
    }

    if (mpiRank == 1) {
      std::cout << "RANKS dofs = ";
      for (int i = 0; i < rankDofs.size(); i++) 
	std::cout << rankDofs[i] << " ";
      std::cout << std::endl;
    }

    // === Remove all macro elements that are not part of the rank partition. ===

    std::vector<MacroElement*> macrosToRemove;
    for (std::deque<MacroElement*>::iterator it = mesh->firstMacroElement();
	 it != mesh->endOfMacroElements();
	 ++it) {
      PartitionElementData *partitionData = 
	dynamic_cast<PartitionElementData*>
	((*it)->getElement()->getElementData(PARTITION_ED));
      if (partitionData->getPartitionStatus() != IN) {
	macrosToRemove.push_back(*it);
      }
    }

    mesh->removeMacroElements(macrosToRemove);

    // === Create local and global dofs ordering. ===

    int *gOrder = (int*)(malloc(sizeof(int) * rankDofs.size()));
    int *lOrder = (int*)(malloc(sizeof(int) * rankDofs.size()));

    for (std::vector<int>::iterator it = rankDofs.begin();
	 it != rankDofs.end(); ++it) {
      gOrder[nRankDOFs++] = *it;
    }

    int rstart = 0;
    MPI_Scan(&nRankDOFs, &rstart, 1, MPI_INT, MPI_SUM, PETSC_COMM_WORLD);
    rstart -= nRankDOFs;
    
    for (int i = 0; i < nRankDOFs; i++) {
      lOrder[i] = rstart + i;
    }

    AOCreateBasic(PETSC_COMM_WORLD, nRankDOFs, gOrder, lOrder, &applicationOrdering);
    
    free(gOrder);
    free(lOrder);
  }

  void ParallelDomainProblemBase::exitParallelization(AdaptInfo *adaptInfo)
  {
    AODestroy(applicationOrdering);
  }

  double ParallelDomainProblemBase::setElemWeights(AdaptInfo *adaptInfo) 
  {
    double localWeightSum = 0.0;
    int elNum = -1;

    elemWeights.clear();

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1,
					 Mesh::CALL_EVERY_EL_PREORDER);
    while (elInfo) {
      Element *element = elInfo->getElement();

      // get partition data
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));

      if (partitionData && partitionData->getPartitionStatus() == IN) {
	if (partitionData->getLevel() == 0) {
	  elNum = element->getIndex();
	}
	TEST_EXIT(elNum != -1)("invalid element number\n");
	if (element->isLeaf()) {
	  elemWeights[elNum] += 1.0;
	  localWeightSum += 1.0;
	}
      }

      elInfo = stack.traverseNext(elInfo);
    }

    return localWeightSum;
  }

  void ParallelDomainProblemBase::partitionMesh(AdaptInfo *adaptInfo)
  {
    if (initialPartitionMesh) {
      initialPartitionMesh = false;
      partitioner->fillCoarsePartitionVec(&oldPartitionVec);
      partitioner->partition(&elemWeights, INITIAL);
    } else {
      oldPartitionVec = partitionVec;
      partitioner->partition(&elemWeights, ADAPTIVE_REPART, 100.0 /*0.000001*/);
    }    

    partitioner->fillCoarsePartitionVec(&partitionVec);
  }

  ParallelDomainProblemScal::ParallelDomainProblemScal(const std::string& name,
						       ProblemScal *problem,
						       ProblemInstatScal *problemInstat)
    : ParallelDomainProblemBase(name, problem, problemInstat, problem->getFESpace())
  {
  }


}