#include "SubAssembler.h"
#include "Assembler.h"
#include "FiniteElemSpace.h"
#include "Operator.h"
#include "BasisFunction.h"
#include "OpenMP.h"
#include "Mesh.h"
#include "Quadrature.h"
#include "DOFVector.h"

namespace AMDiS {

  SubAssembler::SubAssembler(Operator *op,
			     Assembler *assembler,
			     Quadrature *quadrat,
			     int order, 
			     bool optimized,
			     FirstOrderType type) 
    : rowFeSpace(assembler->rowFeSpace),
      colFeSpace(assembler->colFeSpace),
      nRow(0),
      nCol(0),
      coordsAtQPs(NULL),
      coordsNumAllocated(0),
      quadrature(quadrat),
      psiFast(NULL),
      phiFast(NULL),
      symmetric(true),
      opt(optimized),
      firstCall(true),
      name("")
  {
    FUNCNAME("SubAssembler::SubAssembler()");

    TEST_EXIT(rowFeSpace)("No row FE space defined!\n");
    TEST_EXIT(colFeSpace)("No col FE space defined!\n");

    const BasisFunction *psi = rowFeSpace->getBasisFcts();
    const BasisFunction *phi = colFeSpace->getBasisFcts();

    nRow = psi->getNumber();
    nCol = phi->getNumber();

    int maxThreads = omp_get_overall_max_threads();
    terms.resize(maxThreads);

    switch (order) {
    case 0:
      terms = op->zeroOrder;
      break;
    case 1:
      if (type == GRD_PHI)
	terms = op->firstOrderGrdPhi;
      else 
	terms = op->firstOrderGrdPsi;
      break;
    case 2:
      terms = op->secondOrder;
      break;
    }

    // If the number of basis functions in row and col fe space are equal,
    // the element matrix may be symmetric if all operator terms are symmetric.
    // If the row and col fe space are different, the element matrix is neither
    // quadratic, and therefore cannot be symmetric.
    if (nRow == nCol) {
      symmetric = true;
      for (int i = 0; i < static_cast<int>(terms[0].size()); i++) {
	if (!(terms[0][i])->isSymmetric()) {
	  symmetric = false;
	  break;
	}
      }  
    } else {
      symmetric = false;
    }

    dim = rowFeSpace->getMesh()->getDim();
  }


  FastQuadrature *SubAssembler::updateFastQuadrature(FastQuadrature *quadFast,
						     const BasisFunction *psi,
						     Flag updateFlag)
  {
    if (!quadFast) {
      quadFast = FastQuadrature::provideFastQuadrature(psi, *quadrature, updateFlag);
    } else {
      if (!quadFast->initialized(updateFlag))
	quadFast->init(updateFlag);
    }

    return quadFast;
  }


  void SubAssembler::initElement(const ElInfo* smallElInfo,
				 const ElInfo *largeElInfo,
				 Quadrature *quad)
  {
    FUNCNAME("SubAssembler::initElement()");

    // set corrdsAtQPs invalid
    coordsValid = false;

    // set values at QPs invalid
    std::map<const DOFVectorBase<double>*, ValuesAtQPs*>::iterator it1;
    for (it1 = valuesAtQPs.begin(); it1 != valuesAtQPs.end(); ++it1)
      ((*it1).second)->valid = false;

    // set gradients at QPs invalid
    std::map<const DOFVectorBase<double>*, GradientsAtQPs*>::iterator it2;
    for (it2 = gradientsAtQPs.begin(); it2 != gradientsAtQPs.end(); ++it2)
      ((*it2).second)->valid = false;

    int myRank = omp_get_thread_num();
    // calls initElement of each term
    std::vector<OperatorTerm*>::iterator it;
    for (it = terms[myRank].begin(); it != terms[myRank].end(); ++it) {
      if (largeElInfo == NULL)
	(*it)->initElement(smallElInfo, this, quad);
      else
	(*it)->initElement(smallElInfo, largeElInfo, this, quad);      
    }
  }


  WorldVector<double>* SubAssembler::getCoordsAtQPs(const ElInfo* elInfo, 
						    Quadrature *quad)
  {
    Quadrature *localQuad = quad ? quad : quadrature;
  
    const int nPoints = localQuad->getNumPoints();

    // already calculated for this element ?
    if (coordsValid)
      return coordsAtQPs;
   
    if (coordsAtQPs)  {
      if (coordsNumAllocated != nPoints) {
	delete [] coordsAtQPs;
        coordsAtQPs = new WorldVector<double>[nPoints];
	coordsNumAllocated = nPoints;
      }
    } else {
      coordsAtQPs = new WorldVector<double>[nPoints];
      coordsNumAllocated = nPoints;
    }

    // set new values
    WorldVector<double>* k = &(coordsAtQPs[0]);
    for (int l = 0; k < &(coordsAtQPs[nPoints]); ++k, ++l)
      elInfo->coordToWorld(localQuad->getLambda(l), *k);

    // mark values as valid
    coordsValid = true;

    return coordsAtQPs;
  }


  WorldVector<double>* SubAssembler::getGradientsAtQPs(DOFVectorBase<double>* vec, 
						       const ElInfo* elInfo,
						       Quadrature *quad)
  {
    FUNCNAME("SubAssembler::getGradientsAtQPs()");

    TEST_EXIT_DBG(vec)("No DOF vector!\n");

    if (gradientsAtQPs[vec] && gradientsAtQPs[vec]->valid) 
      return &(gradientsAtQPs[vec]->values[0]);

    Quadrature *localQuad = quad ? quad : quadrature;

    if (!gradientsAtQPs[vec])
      gradientsAtQPs[vec] = new GradientsAtQPs;

    gradientsAtQPs[vec]->values.resize(localQuad->getNumPoints());

    WorldVector<double> *values = &(gradientsAtQPs[vec]->values[0]);

    const BasisFunction *psi = rowFeSpace->getBasisFcts();
    const BasisFunction *phi = colFeSpace->getBasisFcts();

    bool sameFeSpaces = 
      vec->getFeSpace() == rowFeSpace || vec->getFeSpace() == colFeSpace;

    if (opt && !quad && sameFeSpaces) {
      if (vec->getFeSpace()->getBasisFcts() == psi)
	psiFast = updateFastQuadrature(psiFast, psi, INIT_GRD_PHI);
      else if(vec->getFeSpace()->getBasisFcts() == phi)
	phiFast = updateFastQuadrature(phiFast, phi, INIT_GRD_PHI);
    }
  
    // calculate new values
    const BasisFunction *basFcts = vec->getFeSpace()->getBasisFcts();

    if (opt && !quad && sameFeSpaces) {
      if (psiFast->getBasisFunctions() == basFcts)
	vec->getGrdAtQPs(elInfo, NULL, psiFast, values);
      else
	vec->getGrdAtQPs(elInfo, NULL, phiFast, values);      
    } else {
      vec->getGrdAtQPs(elInfo, localQuad, NULL, values);
    }
   
    gradientsAtQPs[vec]->valid = true;

    return values;
  }


  WorldVector<double>* SubAssembler::getGradientsAtQPs(DOFVectorBase<double>* vec, 
						       const ElInfo* smallElInfo,
						       const ElInfo* largeElInfo,
						       Quadrature *quad)
  {
    FUNCNAME("SubAssembler::getGradientsAtQPs()");

    TEST_EXIT_DBG(vec)("No DOF vector!\n");

    if (gradientsAtQPs[vec] && gradientsAtQPs[vec]->valid) 
      return &(gradientsAtQPs[vec]->values[0]);

    Quadrature *localQuad = quad ? quad : quadrature;

    if (!gradientsAtQPs[vec])
      gradientsAtQPs[vec] = new GradientsAtQPs;    
    gradientsAtQPs[vec]->values.resize(localQuad->getNumPoints());

    WorldVector<double> *values = &(gradientsAtQPs[vec]->values[0]);

    gradientsAtQPs[vec]->valid = true;
    vec->getGrdAtQPs(smallElInfo, largeElInfo, localQuad, NULL, values);
    return values;
  }

}