#include "DualTraverse.h"
#include "Mesh.h"
#include "ElInfo.h"

namespace AMDiS {

  bool DualTraverse::traverseFirst(Mesh *mesh1, 
				   Mesh *mesh2, 
				   int level1, 
				   int level2, 
				   Flag flag1,
				   Flag flag2,
				   ElInfo **elInfo1,
				   ElInfo **elInfo2,
				   ElInfo **elInfoSmall,
				   ElInfo **elInfoLarge)
  {
    // replace CALL_EL_LEVEL by CALL_MG_LEVEL (covers whole domain)
    if (flag1.isSet(Mesh::CALL_EL_LEVEL)) {
      flag1 &= ~Mesh::CALL_EL_LEVEL;
      flag1 |= Mesh::CALL_MG_LEVEL;
      level1_ = level1;
    } else {
      level1_ = -1;
    }

    if (flag2.isSet(Mesh::CALL_EL_LEVEL)) {
      flag2 &= ~Mesh::CALL_EL_LEVEL;
      flag2 |= Mesh::CALL_MG_LEVEL;
      level2_ = level2;
    } else {
      level2_ = -1;
    }

    // replace CALL_LEAF_EL_LEVEL by CALL_MG_LEVEL (covers whole domain)
    if (flag1.isSet(Mesh::CALL_LEAF_EL_LEVEL)) {
      flag1 &= ~Mesh::CALL_LEAF_EL_LEVEL;
      flag1 |= Mesh::CALL_MG_LEVEL;
      level1_ = level1;
      callLeafElLevel1_ = true;
    } else {
      level1_ = -1;
      callLeafElLevel1_ = false;
    }

    if (flag2.isSet(Mesh::CALL_LEAF_EL_LEVEL)) {
      flag2 &= ~Mesh::CALL_LEAF_EL_LEVEL;
      flag2 |= Mesh::CALL_MG_LEVEL;
      level2_ = level2;
      callLeafElLevel2_ = true;
    } else {
      level2_ = -1;
      callLeafElLevel2_ = false;
    }

    // call standard traverse
    *elInfo1 = stack1.traverseFirst(mesh1, level1, flag1);
    *elInfo2 = stack2.traverseFirst(mesh2, level2, flag2);
 
    // finished ?
    if (*elInfo1 == NULL || *elInfo2 == NULL) {
      TEST_EXIT(*elInfo1 == *elInfo2)("invalid dual traverse\n");
      return false;
    }

    rest = 1.0;

    bool accepted = check(elInfo1, elInfo2, elInfoSmall, elInfoLarge);

    // check for non domain covering level traverse
    if (!accepted ||
	(level1_ != -1 && (*elInfo1)->getLevel() != level1_) ||
	(callLeafElLevel1_ && !(*elInfo1)->getElement()->isLeaf()) ||
	(level2_ != -1 && (*elInfo2)->getLevel() != level2_) ||
	(callLeafElLevel2_ && !(*elInfo2)->getElement()->isLeaf())) {
      return traverseNext(elInfo1, elInfo2, elInfoSmall, elInfoLarge);
    }

    fillSubElInfo(*elInfo1, *elInfo2, *elInfoSmall, *elInfoLarge);

    return true;
  }

  bool DualTraverse::traverseNext(ElInfo **elInfo1,
				  ElInfo **elInfo2,
				  ElInfo **elInfoSmall,
				  ElInfo **elInfoLarge)
  {
    // call standard traverse
    if (inc1) {
      *elInfo1 = stack1.traverseNext(*elInfo1);
    }
    if (inc2) {
      *elInfo2 = stack2.traverseNext(*elInfo2);
    }

    // finished ?
    if (*elInfo1 == NULL || *elInfo2 == NULL) {
      TEST_EXIT(*elInfo1 == *elInfo2)("invalid dual traverse\n");
      return false;
    }

    bool accepted = check(elInfo1, elInfo2, elInfoSmall, elInfoLarge);

    // check for non domain covering level traverse
    if (!accepted ||
	(level1_ != -1 && (*elInfo1)->getLevel() != level1_) ||
	(callLeafElLevel1_ && !(*elInfo1)->getElement()->isLeaf()) ||
	(level2_ != -1 && (*elInfo2)->getLevel() != level2_) ||
	(callLeafElLevel2_ && !(*elInfo2)->getElement()->isLeaf())) {
      return traverseNext(elInfo1, elInfo2, elInfoSmall, elInfoLarge);
    }

    fillSubElInfo(*elInfo1, *elInfo2, *elInfoSmall, *elInfoLarge);

    return true;
  }

  void DualTraverse::prepareNextStep(ElInfo **elInfo1,
				     ElInfo **elInfo2,
				     ElInfo **elInfoSmall,
				     ElInfo **elInfoLarge)
  {
    // which is the smaller element ?
    *elInfoSmall = 
      (*elInfo1)->getLevel() > (*elInfo2)->getLevel() ?
      *elInfo1 :
      *elInfo2;
    *elInfoLarge = 
      (*elInfo1)->getLevel() <= (*elInfo2)->getLevel() ?
      *elInfo1 :
      *elInfo2;

    // update rest
    double newRest = rest - 1.0 / (1 << ((*elInfoSmall)->getLevel() - (*elInfoLarge)->getLevel()));

    if (newRest < 1e-32) {
      // large element finished -> increment both elements
      rest = 1.0; 
      inc1 = true;
      inc2 = true;
      stack1.stopDisplacementCalculation();
      stack2.stopDisplacementCalculation();
    } else {
      // increment only small element
      inc1 = (*elInfo1 == *elInfoSmall) ? true : false;
      inc2 = (*elInfo2 == *elInfoSmall) ? true : false;

      // If rest is 1, than this is the first element pair with this large
      // element. Therefore, fill the displacement stack of the small element
      // traverse stack.
      if (rest == 1.0) {
	if (*elInfo1 == *elInfoSmall) {
	  stack1.startDisplacementCalculation((*elInfo2)->getLevel());
	  //	  (*elInfo1)->setDisplacement(0);
	} else {
	  stack2.startDisplacementCalculation((*elInfo1)->getLevel());
	  //	  (*elInfo2)->setDisplacement(0);
	}
      }

      rest = newRest;
    }
  }

  void DualTraverse::fillSubElInfo(ElInfo *elInfo1, 
				   ElInfo *elInfo2,
				   ElInfo *elInfoSmall,
				   ElInfo *elInfoLarge)
  {
    if (!fillSubElemMat)
      return;

    VectorOfFixVecs<DimVec<double> > *subCoords = 
      NEW VectorOfFixVecs<DimVec<double> >(elInfo1->getMesh()->getDim(), 
					   elInfo1->getMesh()->getDim() + 1, 
					   NO_INIT);
    
    DimMat<double> *subCoordsMat = elInfoSmall->getSubElemCoordsMat();
    if (!subCoordsMat) {
      subCoordsMat = NEW DimMat<double>(elInfo1->getMesh()->getDim());
    }

    if (elInfo1 == elInfoSmall) {
      stack1.getCoordsInElem(elInfo2, subCoords);
    } else {
      stack2.getCoordsInElem(elInfo1, subCoords);
    }

    for (int i = 0; i < elInfo1->getMesh()->getDim() + 1; i++) {
      for (int j = 0; j < elInfo1->getMesh()->getDim() + 1; j++) {
	(*subCoordsMat)[j][i] = (*subCoords)[i][j];
      }
    }
       
    elInfoSmall->setSubElemCoordsMat(subCoordsMat);

    DELETE subCoords;
  }
}