//
// 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 "parallel/CheckerPartitioner.h"
#include "Traverse.h"


namespace AMDiS {
  
  void CheckerPartitioner::createInitialPartitioning()
  {
    FUNCNAME("CheckerPartitioner::createInitialPartitioning()");

    // In one of the stripes mode, we have to check if the number of macro
    // elements with together with the number of nodes.
    if (mode == 1 || mode == 2 || mode == 3) {
      int elCounter = 0;
      TraverseStack stack;
      ElInfo *elInfo = stack.traverseFirst(mesh, 0, Mesh::CALL_EL_LEVEL);
      while (elInfo) {
	elCounter++;
	elInfo = stack.traverseNext(elInfo);
      }
      
      if (mesh->getDim() == 2)
	TEST_EXIT(elCounter == 2 * mpiSize * mpiSize)
	  ("The number of macro elements is %d, but must be %d for %d number of nodes!",
	   elCounter, 2 * mpiSize * mpiSize, mpiSize);

      if (mesh->getDim() == 3)
	TEST_EXIT(elCounter == 6 * static_cast<int>(pow(mpiSize, 1.5)))
	  ("The number of macro elements is %d, but must be %d for %d number of nodes!",
	   elCounter, 6 * static_cast<int>(pow(mpiSize, 1.5)), mpiSize);

    }

    if (multilevel) {
      TEST_EXIT(MPI::COMM_WORLD.Get_size() == 16)
	("Multilevel partitioning is implemented for 16 nodes only!\n");
    }
    
    int dim = mesh->getDim();
    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, 0, Mesh::CALL_EL_LEVEL);
    while (elInfo) {
      Element *el = elInfo->getElement();
      int elIndex = el->getIndex();
      int boxIndex = elIndex / (dim == 2 ? 2 : 6);
      int elInRank = -1;

      switch (mode) {
      case 0:
	if (!multilevel) {
	  elInRank = boxIndex;	
	} else {
	  TEST_EXIT_DBG(boxIndex >= 0 && boxIndex <= 15)("Wrong box index!\n");
	  int rs[16] = {0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15};
	  elInRank = rs[boxIndex];
	}
	break;

      case 1:
	// x-slices

	{
	  if (dim == 2)
	    elInRank = elIndex / (2 * mpiSize);
	  
	  if (dim == 3) {
	    int boxSliceY = 
	      (boxIndex % mpiSize) / static_cast<int>(sqrt(mpiSize));
	    int boxSliceZ = boxIndex / mpiSize;
	    elInRank = boxSliceY * static_cast<int>(sqrt(mpiSize)) + boxSliceZ;
	  }
	}

	break;

      case 2:
	// y-slices

	{
	  if (dim == 2)
	    elInRank = (elIndex % (2 * mpiSize)) / 2;
	  
	  if (dim == 3) {
	    int boxSliceX = 
	      (boxIndex % mpiSize) % static_cast<int>(sqrt(mpiSize));
	    int boxSliceZ = boxIndex / mpiSize;
	    elInRank = boxSliceX * static_cast<int>(sqrt(mpiSize)) + boxSliceZ;
	  }
	}

	break;

      case 3:
	// z-slices

	{
	  int boxSliceX = (boxIndex % mpiSize) % static_cast<int>(sqrt(mpiSize));
	  int boxSliceY = (boxIndex % mpiSize) / static_cast<int>(sqrt(mpiSize));

	  elInRank = boxSliceX * static_cast<int>(sqrt(mpiSize)) + boxSliceY;
	}

	break;

      default:
	ERROR_EXIT("Mode %d does not exists for checker based mesh partitioning!\n", 
		   mode);
      }

      TEST_EXIT_DBG(elInRank >= 0)("Should not happen!\n");
      TEST_EXIT_DBG(elInRank < mpiSize)("Should not happen!\n");
      
      elementInRank[elIndex] = (elInRank == mpiRank);
      partitionMap[elIndex] = elInRank;	
      
      elInfo = stack.traverseNext(elInfo);
    }
  }

}