#include "SMIAdapter.h"
#include "Mesh.h"
#include "Traverse.h"
#include "ElInfo.h"
#include "Flag.h"
#include "FiniteElemSpace.h"
#include "BasisFunction.h"
#include "DOFVector.h"
#include "DOFAdmin.h"
#include "ElementRegion_ED.h"
#include "SurfaceRegion_ED.h"
#include "Flag.h"

namespace AMDiS {

  SMIAdapter::SMIAdapter(int              smiApplicationId,
			 int              smiMeshId,
			 FiniteElemSpace *feSpace,
			 int              elementRegion,
			 int              surfaceRegion,
			 bool (*elementFct)(ElInfo *elInfo),
			 bool (*surfaceFct)(ElInfo *elInfo, int side)) 
    : smiApplicationId_(smiApplicationId), 
      smiMeshId_(smiMeshId),
      feSpace_(feSpace),
      elementRegion_(elementRegion),
      surfaceRegion_(surfaceRegion),
      addNeighbourInfo_(false),
      newNodeIndex_(NULL),
      oldNodeIndex_(NULL),
      newElementIndex_(NULL),
      oldElementIndex_(NULL),
      elementFct_(elementFct),
      surfaceFct_(surfaceFct)
  {
    TEST_EXIT(!surfaceFct_)("surfaceFct not yet supported\n");
    TEST_EXIT(elementFct_ == NULL || elementRegion_ == -1)
      ("don't use elementRegion AND elementFct\n");

    Mesh *mesh = feSpace_->getMesh();
    int dim = mesh->getDim();

    if(surfaceRegion > -1) {
      dim -= 1;
    }

    TEST_EXIT(dim > 0 && dim < 4)("invalid dim\n");

    elementType_ = dim * 10 + 1; // linear elements only

    int smiError = SMI_Add_application(smiApplicationId_,
				       SMI_APP_MODE_DEFAULT);

    smiError = SMI_Add_mesh(smiApplicationId_,
			    smiMeshId_, SMI_MESH_MODE_DEFAULT);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Add_mesh() failed with error %d\n", smiError);

    smiError = SMI_Begin_write_transaction(smiApplicationId_,
					   smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Begin_modification() failed with error %d\n", smiError);

    smiError = SMI_Add_elem_type(smiApplicationId_,
				 smiMeshId_,
				 elementType_,
				 dim + 1);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Add_elem_type() failed with error %d\n", smiError);

    smiError = SMI_End_write_transaction(smiApplicationId_,
					 smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_End_modification() failed with error %d\n", smiError);
  }

  void SMIAdapter::addQuantity(int quantityId, DOFVector<double> *dofVector) {
    TEST_EXIT(dofVectors_[quantityId].size() == 0)
      ("quantityId already added\n");

    dofVectors_[quantityId].push_back(dofVector);

    static double defaultValue = 0.0;

    int smiError;

    smiError = SMI_Begin_write_transaction(smiApplicationId_,
					   smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Begin_modification() failed with error %d\n", smiError);

    smiError = SMI_Add_quantity(smiApplicationId_,
				smiMeshId_,
				quantityId,
				SMI_LOC_NODE,
				SMI_TYPE_DOUBLE,
				1,
				&defaultValue);
      
    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Add_quantity faild with error %d\n", smiError);

    smiError = SMI_End_write_transaction(smiApplicationId_,
					 smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_End_modification() failed with error %d\n", smiError);
  }

  void SMIAdapter::addQuantity(int quantityId,
			       int quantityDim,
			       DOFVector<double> **dofVector) 
  {
    TEST_EXIT(dofVectors_[quantityId].size() == 0)
      ("quantityId already added\n");

    dofVectors_[quantityId].resize(quantityDim);

    int comp;
    for(comp = 0; comp < quantityDim; comp++) {
      dofVectors_[quantityId][comp] = dofVector[comp];
    }

    TEST_EXIT(quantityDim < 4)("quantityDim > 3\n");

    static double defaultValue[3] = {0.0, 0.0, 0.0};

    int smiError;

    smiError = SMI_Begin_write_transaction(smiApplicationId_,
					   smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Begin_modification() failed with error %d\n", smiError);

    smiError = SMI_Add_quantity(smiApplicationId_,
				smiMeshId_,
				quantityId,
				SMI_LOC_NODE,
				SMI_TYPE_DOUBLE,
				quantityDim,
				defaultValue);
      
    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Add_quantity faild with error %d\n", smiError);

    smiError = SMI_End_write_transaction(smiApplicationId_,
					 smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_End_modification() failed with error %d\n", smiError);
  }

  void SMIAdapter::transferMeshToSMI()
  {
    int smiError = SMI_Begin_write_transaction(smiApplicationId_,
					       smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Begin_modification() failed with error %d\n", smiError);

    smiError = SMI_Clear_mesh(smiApplicationId_,
			      smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Clear_mesh() failed with error %d\n", smiError);

    int i, j;
    Mesh *mesh = feSpace_->getMesh();
    int dim = mesh->getDim();
    DOFAdmin *admin = feSpace_->getAdmin();
    const BasisFunction *basFcts = feSpace_->getBasisFcts();
    int numBasFcts = basFcts->getNumber();

    Element *element;
    ElementData *elementData;
    ElementData *regionData;
    int region;
    int side;
    int numNewNodes;
    int numNodes;
    bool validElement;
    double *nodeCoords = GET_MEMORY(double, numBasFcts * dim); 

    DOFVector<char> alreadyAdded(feSpace_, "already added nodes");
    alreadyAdded.set(0);

    int elementCounter[1] = { 0 };
    int elementIndex;

    DimVec<double> *bary = NULL;
    WorldVector<double> world;

    const DegreeOfFreedom *elementDofs = NULL;
    DegreeOfFreedom dof;
    DegreeOfFreedom *nodeIndices = GET_MEMORY(int, numBasFcts);
    DegreeOfFreedom *newNodeIndices = GET_MEMORY(int, numBasFcts);

    Flag fillFlag = Mesh::CALL_LEAF_EL | Mesh::FILL_COORDS;
  
    if(addNeighbourInfo_) {
      fillFlag |= Mesh::FILL_NEIGH;
    }

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, fillFlag);

    while(elInfo) {
    
      element = elInfo->getElement();
      elementData = element->getElementData();
      elementDofs = basFcts->getLocalIndices(element, admin, NULL);

      // check element region
      if(elementRegion_ > -1) {
	regionData = elementData->getElementData(ELEMENT_REGION);
	if(regionData) {
	  region = dynamic_cast<ElementRegion_ED*>(regionData)->getRegion();
	} else {
	  region = -1;
	}
	validElement = (region == elementRegion_);
      } else {
	validElement = true;
      }

      if(elementFct_) {
	validElement = (*elementFct_)(elInfo);
      }
    
      if(validElement) {

	// check surface region
	if(surfaceRegion_ > -1) {
	  regionData = elementData->getElementData(SURFACE_REGION);
	  while(regionData) {
	    region = dynamic_cast<SurfaceRegion_ED*>(regionData)->getRegion();

	    // surface region found ?
	    if(region == surfaceRegion_) {
	      // add surface element
	      side = dynamic_cast<SurfaceRegion_ED*>(regionData)->getSide();
	      numNewNodes = 0;
	      numNodes = 0;
	      for(i = 0; i < numBasFcts; i++) {
		bary = basFcts->getCoords(i);
		if((*bary)[side] == 0) {
		  dof = elementDofs[i];

		  if(newNodeIndex_) {
		    dof = (*(newNodeIndex_))[dof] - 1;
		    TEST_EXIT(dof >= 0)("invalid node index\n");
		  }
		
		  if(!alreadyAdded[dof]) {
		    newNodeIndices[numNewNodes] = dof;
		    elInfo->coordToWorld(*bary, world);
		    for(j = 0; j < dim; j++) {
		      nodeCoords[numNewNodes * dim + j] = world[j];
		    }
		    numNewNodes++;
		    alreadyAdded[dof] = 1;
		  }
		  nodeIndices[numNodes] = dof;
		  numNodes++;
		}
	      }

	      smiError = SMI_Add_nodes(smiApplicationId_,
				       smiMeshId_,
				       numNewNodes,
				       newNodeIndices,
				       nodeCoords,
				       dim);

	      TEST_EXIT(smiError == SMI_OK)
		("SMI_Add_nodes faild with error %d\n", smiError);

	      smiError = SMI_Add_elems(smiApplicationId_,
				       smiMeshId_,
				       1,
				       &elementType_,
				       dim+1,
				       nodeIndices,
				       elementCounter);

	      (elementCounter[0])++;

	      TEST_EXIT(smiError == SMI_OK)
		("SMI_Add_elems faild with error %d\n", smiError);
	    }
	    regionData = regionData->getDecorated(SURFACE_REGION);
	  }
	} else { 
	  // add volume element


	  numNewNodes = 0;
	  for(i = 0; i < numBasFcts; i++) {
	    bary = basFcts->getCoords(i);
	    dof = elementDofs[i];

	    if(newNodeIndex_) {
	      dof = (*(newNodeIndex_))[dof] - 1;
	      TEST_EXIT(dof >= 0)("invalid node index\n");
	    }

	    if(!alreadyAdded[dof]) {
	      newNodeIndices[numNewNodes] = dof;
	      elInfo->coordToWorld(*bary, world);
	      for(j = 0; j < dim; j++) {
		nodeCoords[numNewNodes * dim + j] = world[j];
	      }
	      numNewNodes++;
	      alreadyAdded[dof] = 1;
	    }
	  }

	  smiError = SMI_Add_nodes(smiApplicationId_,
				   smiMeshId_,
				   numNewNodes,
				   newNodeIndices,
				   nodeCoords,
				   dim);

	  TEST_EXIT(smiError == SMI_OK)
	    ("SMI_Add_nodes faild with error %d\n", smiError);

	  elementIndex = element->getIndex();

	  if(newElementIndex_) {
	    elementIndex = (*(newElementIndex_))[elementIndex] - 1;
	    TEST_EXIT(elementIndex >= 0)("invalid element index\n");
	  }

	  smiError = SMI_Add_elems(smiApplicationId_,
				   smiMeshId_,
				   1,
				   &elementType_,
				   dim+1,
				   const_cast<int*>(elementDofs),
				   &elementIndex);

	  if(addNeighbourInfo_) {
	    DimVec<int> neighbours(dim, NO_INIT);
	    DimVec<int> oppVertices(dim, NO_INIT);
	    for(i = 0; i < dim + 1; i++) {
	      if(elInfo->getNeighbour(i)) {
		neighbours[i] = elInfo->getNeighbour(i)->getIndex();
		if(newElementIndex_) {
		  neighbours[i] = (*(newElementIndex_))[neighbours[i]] - 1;
		  TEST_EXIT(neighbours[i] >= 0)("invalid element index\n");
		}
		oppVertices[i] = elInfo->getOppVertex(i);
	      } else {
		neighbours[i] = -1;
		oppVertices[i] = -1;
	      }
	    }
	    SMI_Set_quantity_values(smiApplicationId_,
				    smiMeshId_,
				    SMI_NEIGH_QUANTITY,
				    SMI_TYPE_INT,
				    dim+1,
				    1,
				    &elementIndex,
				    neighbours.getValArray());
	    SMI_Set_quantity_values(smiApplicationId_,
				    smiMeshId_,
				    SMI_OPP_V_QUANTITY,
				    SMI_TYPE_INT,
				    dim+1,
				    1,
				    &elementIndex,
				    oppVertices.getValArray());
	  }


	  (elementCounter[0])++;
	
	  TEST_EXIT(smiError == SMI_OK)
	    ("SMI_Add_elems faild with error %d\n", smiError);
	}
      }
    
      elInfo = stack.traverseNext(elInfo);
    }
  
    FREE_MEMORY(nodeCoords, double, numBasFcts * dim);
    FREE_MEMORY(nodeIndices, int, numBasFcts);
    FREE_MEMORY(newNodeIndices, int, numBasFcts);

    smiError = SMI_End_write_transaction(smiApplicationId_,
					 smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_End_modification() failed with error %d\n", smiError);
  }

  void SMIAdapter::transferQuantitiesToSMI(int quantityID)
  {
    int smiError;
  
    smiError = SMI_Begin_write_transaction(smiApplicationId_,
					   smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Begin_modification() failed with error %d\n", smiError);

    int  numNodes;
    int *nodeIndices;

    smiError = SMI_Get_all_nodes(smiApplicationId_,
				 smiMeshId_,
				 &numNodes,
				 &nodeIndices);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Get_all_nodes() failed with error %d\n", smiError);

    //double *values = GET_MEMORY(double, numNodes);

    if(quantityID == -1) {
      std::map<int, std::vector<DOFVector<double>*> >::iterator quantityIt;
      std::map<int, std::vector<DOFVector<double>*> >::iterator quantityEnd = 
	  dofVectors_.end();

      for(quantityIt = dofVectors_.begin(); quantityIt != quantityEnd; ++quantityIt) {
	int quantityDim = static_cast<int>(quantityIt->second.size());

	double *values = GET_MEMORY(double, numNodes * quantityDim);

	int i, comp;
	for(i = 0; i < numNodes; i++) {
	  DegreeOfFreedom index = nodeIndices[i];
	  if(oldNodeIndex_) {
	    index = (*(oldNodeIndex_))[index] - 1;
	  }
	  for(comp = 0; comp < quantityDim; comp++) {
	    DOFVector<double> *dofVector = (quantityIt->second)[comp];
	    values[i * quantityDim + comp] = (*dofVector)[index];
	  }
	}

	smiError = SMI_Set_quantity_values(smiApplicationId_,
					   smiMeshId_,
					   (*quantityIt).first,
					   SMI_TYPE_DOUBLE,
					   quantityDim,
					   numNodes,
					   nodeIndices,
					   values);

	TEST_EXIT(smiError == SMI_OK)
	  ("SMI_Set_quantity_values() failed with error %d\n", smiError);
      
	FREE_MEMORY(values, double, numNodes * quantityDim);
      }
    } else {
      int quantityDim = static_cast<int>(dofVectors_[quantityID].size());

      double *values = GET_MEMORY(double, numNodes * quantityDim);

      int i, comp;
      for(i = 0; i < numNodes; i++) {
	DegreeOfFreedom index = nodeIndices[i];
	if(oldNodeIndex_) {
	  index = (*(oldNodeIndex_))[index] - 1;
	}
	for(comp = 0; comp < quantityDim; comp++) {
	  DOFVector<double> *dofVector = dofVectors_[quantityID][comp];
	  values[i*quantityDim + comp] = (*dofVector)[index];
	}
      }
	
      smiError = SMI_Set_quantity_values(smiApplicationId_,
					 smiMeshId_,
					 quantityID,
					 SMI_TYPE_DOUBLE,
					 quantityDim,
					 numNodes,
					 nodeIndices,
					 values);

      TEST_EXIT(smiError == SMI_OK)
	("SMI_Set_quantity_values() failed with error %d\n", smiError);

      FREE_MEMORY(values, double, numNodes * quantityDim);
    }

    //FREE_MEMORY(values, double, numNodes);

    smiError = SMI_End_write_transaction(smiApplicationId_,
					 smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_End_modification() failed with error %d\n", smiError);
  }

  void SMIAdapter::getQuantitiesFromSMI(int quantityID)
  {
    int smiError;
  
    smiError = SMI_Begin_read_transaction(smiApplicationId_,
					  smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Begin_modification() failed with error %d\n", smiError);

    int  numNodes;
    int *nodeIndices;

    smiError = SMI_Get_all_nodes(smiApplicationId_,
				 smiMeshId_,
				 &numNodes,
				 &nodeIndices);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_Get_all_nodes() failed with error %d\n", smiError);

    //double *values = GET_MEMORY(double, numNodes);

    if(quantityID == -1) {
      std::map<int, std::vector<DOFVector<double>*> >::iterator quantityIt;
      std::map<int, std::vector<DOFVector<double>*> >::iterator quantityEnd = 
	  dofVectors_.end();
      for(quantityIt = dofVectors_.begin(); quantityIt != quantityEnd; ++quantityIt) {
	int quantityDim = static_cast<int>(quantityIt->second.size());

	double *values = GET_MEMORY(double, numNodes * quantityDim);

	smiError = SMI_Get_quantity_values(smiApplicationId_,
					   smiMeshId_,
					   (*quantityIt).first,
					   SMI_TYPE_DOUBLE,
					   quantityDim,
					   numNodes,
					   nodeIndices,
					   values);

	TEST_EXIT(smiError == SMI_OK)
	  ("SMI_Set_quantity_values() failed with error %d\n", smiError);

	int i, comp;
	for(i = 0; i < numNodes; i++) {
	  DegreeOfFreedom index = nodeIndices[i];
	  if(oldNodeIndex_) {
	    index = (*(oldNodeIndex_))[index] - 1;
	  }
	  for(comp = 0; comp < quantityDim; comp++) {
	    DOFVector<double> *dofVector = (quantityIt->second)[comp];
	    (*dofVector)[index] = values[i * quantityDim + comp];
	  }
	}

	FREE_MEMORY(values, double, numNodes * quantityDim);
      }
    } else {
      int quantityDim = static_cast<int>(dofVectors_[quantityID].size());

      double *values = GET_MEMORY(double, numNodes * quantityDim);

      smiError = SMI_Get_quantity_values(smiApplicationId_,
					 smiMeshId_,
					 quantityID,
					 SMI_TYPE_DOUBLE,
					 quantityDim,
					 numNodes,
					 nodeIndices,
					 values);

      TEST_EXIT(smiError == SMI_OK)
	("SMI_Set_quantity_values() failed with error %d\n", smiError);
    
      int i, comp;
      for(i = 0; i < numNodes; i++) {
	DegreeOfFreedom index = nodeIndices[i];
	if(oldNodeIndex_) {
	  index = (*(oldNodeIndex_))[index] - 1;
	}
	for(comp = 0; comp < quantityDim; comp++) {
	  DOFVector<double> *dofVector = dofVectors_[quantityID][comp];
	  (*dofVector)[index] = values[i * quantityDim + comp];
	}
      }

      FREE_MEMORY(values, double, numNodes * quantityDim);
    }

    //FREE_MEMORY(values, double, numNodes);

    smiError = SMI_End_read_transaction(smiApplicationId_,
					smiMeshId_);

    TEST_EXIT(smiError == SMI_OK)
      ("SMI_End_modification() failed with error %d\n", smiError);
  }

  void SMIAdapter::addNeighbourInfo() 
  {
    TEST_EXIT(surfaceRegion_ == -1)
      ("no neighbour info available for surface meshes\n");

    int dim = feSpace_->getMesh()->getDim();

    if(SMI_OK != SMI_Get_quantity_info(smiApplicationId_,
				       smiMeshId_,
				       SMI_NEIGH_QUANTITY,
				       NULL, NULL, NULL))
      {
	int defaultValue[3] = {-1, -1, -1};

	SMI_Begin_write_transaction(smiApplicationId_,
				    smiMeshId_);

	SMI_Add_quantity(smiApplicationId_,
			 smiMeshId_,
			 SMI_NEIGH_QUANTITY,
			 SMI_LOC_ELEM,
			 SMI_TYPE_INT,
			 dim+1,
			 defaultValue);

	SMI_Add_quantity(smiApplicationId_,
			 smiMeshId_,
			 SMI_OPP_V_QUANTITY,
			 SMI_LOC_ELEM,
			 SMI_TYPE_INT,
			 dim+1,
			 defaultValue);

	SMI_End_write_transaction(smiApplicationId_,
				  smiMeshId_);
      }

    addNeighbourInfo_ = true;
  }

  void SMIAdapter::getNeighbourInfo(int  elementIndex,
				    int *neighbours,
				    int *oppVertices)
  {
    int i, dim = feSpace_->getMesh()->getDim();

    TEST_EXIT(addNeighbourInfo_)("no neighbour info added\n");
    SMI_Get_quantity_values(smiApplicationId_,
			    smiMeshId_,
			    SMI_NEIGH_QUANTITY,
			    SMI_TYPE_INT,
			    dim+1,
			    1,
			    &elementIndex,
			    neighbours);
    SMI_Get_quantity_values(smiApplicationId_,
			    smiMeshId_,
			    SMI_OPP_V_QUANTITY,
			    SMI_TYPE_INT,
			    dim+1,
			    1,
			    &elementIndex,
			    oppVertices);

    if(oldElementIndex_) {
      for(i = 0; i < dim + 1; i++) {
	neighbours[i] = (*(oldElementIndex_))[neighbours[i]] - 1;
      }
    }
  }

}