#include <vector>
#include <boost/numeric/mtl/mtl.hpp>

#include "Assembler.h"
#include "ZeroOrderAssembler.h"
#include "Operator.h"
#include "QPsiPhi.h"
#include "FiniteElemSpace.h"
#include "Quadrature.h"
#include "DOFVector.h"
#include "ElementMatrix.h"
#include "OpenMP.h"

namespace AMDiS {

  std::vector<SubAssembler*> ZeroOrderAssembler::optimizedSubAssemblers;
  std::vector<SubAssembler*> ZeroOrderAssembler::standardSubAssemblers;

  ZeroOrderAssembler::ZeroOrderAssembler(Operator *op,
					 Assembler *assembler,
					 Quadrature *quad,
					 bool optimized)
    : SubAssembler(op, assembler, quad, 0, optimized)
  {}

  ZeroOrderAssembler* ZeroOrderAssembler::getSubAssembler(Operator* op,
							  Assembler *assembler,
							  Quadrature *quad,
							  bool optimized)
  {
    // check if a assembler is needed at all
    if (op->zeroOrder.size() == 0)
      return NULL;   

    ZeroOrderAssembler *newAssembler;

    std::vector<SubAssembler*> *subAssemblers =
      optimized ? &optimizedSubAssemblers : &standardSubAssemblers;

    int myRank = omp_get_thread_num();
    std::vector<OperatorTerm*> opTerms = op->zeroOrder[myRank];

    sort(opTerms.begin(), opTerms.end());

    // check if a new assembler is needed
    if (quad) {
      for (int i = 0; i < static_cast<int>( subAssemblers->size()); i++) {
	std::vector<OperatorTerm*> assTerms = *((*subAssemblers)[i]->getTerms());

	sort(assTerms.begin(), assTerms.end());

	if ((opTerms == assTerms) && 
	    ((*subAssemblers)[i]->getQuadrature() == quad)) {
	  return dynamic_cast<ZeroOrderAssembler*>((*subAssemblers)[i]);
	}
      }
    }
 
    // check if all terms are pw_const
    bool pwConst = true;
    for (int i = 0; i < static_cast<int>( op->zeroOrder[myRank].size()); i++) {
      if (!op->zeroOrder[myRank][i]->isPWConst()) {
	pwConst = false;
	break;
      }
    }  

    // create new assembler
    if (!optimized) {
	newAssembler = NEW StandardZOA(op, assembler, quad);
    } else {
      if (pwConst) {
	  newAssembler = NEW PrecalcZOA(op, assembler, quad);
      } else {
	  newAssembler = NEW FastQuadZOA(op, assembler, quad);
      }
    }

    subAssemblers->push_back(newAssembler);
    return newAssembler;
  }

  StandardZOA::StandardZOA(Operator *op, Assembler *assembler, Quadrature *quad)
    : ZeroOrderAssembler(op, assembler, quad, false)
  {}

  void StandardZOA::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix *mat)
  {
    const BasisFunction *psi = owner->getRowFESpace()->getBasisFcts();
    const BasisFunction *phi = owner->getColFESpace()->getBasisFcts();

    int nPoints = quadrature->getNumPoints();
    std::vector<double> c(nPoints, 0.0);
    std::vector<double> phival(nCol);

    int myRank = omp_get_thread_num();
    std::vector<OperatorTerm*>::iterator termIt;
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt) {
      (static_cast<ZeroOrderTerm*>((*termIt)))->getC(elInfo, nPoints, c);
    }
      
    if (symmetric) {
      TEST_EXIT_DBG(nCol == nRow)("nCol != nRow, but symmetric assembling!\n");
            
      for (int iq = 0; iq < nPoints; iq++) {
	c[iq] *= elInfo->getDet();

	// calculate phi at QPs only once!
	for (int i = 0; i < nCol; i++) {
	  phival[i] = (*(phi->getPhi(i)))(quadrature->getLambda(iq));
	}

	for (int i = 0; i < nRow; i++) {
	  double psival = (*(psi->getPhi(i)))(quadrature->getLambda(iq));
	  (*mat)[i][i] += quadrature->getWeight(iq) * c[iq] * psival * phival[i];
	  for (int j = i + 1; j < nCol; j++) {
	    double val = quadrature->getWeight(iq) * c[iq] * psival * phival[j];
	    (*mat)[i][j] += val;
	    (*mat)[j][i] += val;
	  }
	}
      }
    } else {      //  non symmetric assembling 
      for (int iq = 0; iq < nPoints; iq++) {
	c[iq] *= elInfo->getDet();

	// calculate phi at QPs only once!
	for (int i = 0; i < nCol; i++) {
	  phival[i] = (*(phi->getPhi(i)))(quadrature->getLambda(iq));
	}

	for (int i = 0; i < nRow; i++) {
	  double psival = (*(psi->getPhi(i)))(quadrature->getLambda(iq));
	  for (int j = 0; j < nCol; j++) {
	    (*mat)[i][j] += quadrature->getWeight(iq) * c[iq] * psival * phival[j];
	  }
	}
      }
    }
  }

  void StandardZOA::calculateElementMatrix(const ElInfo *rowElInfo,
					   const ElInfo *colElInfo,
					   const ElInfo *smallElInfo,
					   const ElInfo *largeElInfo,
					   ElementMatrix *mat) 
  {
    FUNCNAME("StandardZOA::calculateElementMatrix()");

    TEST_EXIT((nRow <= 3) && (nCol <= 3))("not yet!\n");

    calculateElementMatrix(smallElInfo, mat);

    const mtl::dense2D<double>& m = smallElInfo->getSubElemCoordsMat();

    ElementMatrix tmpMat(nRow, nRow);

    for (int i = 0; i < nRow; i++) {
      for (int j = 0; j < nRow; j++) {
	tmpMat[i][j] = 0.0;

	for (int k = 0; k < nRow; k++) {
	  tmpMat[i][j] += m[i][j] * (*mat)[j][i];
	}
      }
    }

    for (int i = 0; i < nRow; i++) {
      for (int j = 0; j < nRow; j++) {
	(*mat)[i][j] = tmpMat[i][j];
      }
    }

    /*
    const BasisFunction *psi = owner->getRowFESpace()->getBasisFcts();
    const BasisFunction *phi = owner->getColFESpace()->getBasisFcts();
    const mtl::dense2D<double>& m = smallElInfo->getSubElemCoordsMat();

    int myRank = omp_get_thread_num();
    int nPoints = quadrature->getNumPoints();
    std::vector<double> c(nPoints, 0.0);

    std::vector<OperatorTerm*>::iterator termIt;
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt)
      (static_cast<ZeroOrderTerm*>((*termIt)))->getC(smallElInfo, nPoints, c);

    for (int i = 0; i < nRow; i++) {
      for (int j = 0; j < nCol; j++) {

	for (int iq = 0; iq < nPoints; iq++) {
	  double val0 = c[iq] * smallElInfo->getDet() * quadrature->getWeight(iq) * 
	    (*(phi->getPhi(i)))(quadrature->getLambda(iq));

	  double val1 = 0.0;
	  for (int k = 0; k < nCol; k++)
	    val1 += m[j][k] * (*(psi->getPhi(k)))(quadrature->getLambda(iq));

	  val0 *= val1;

	  (*mat)[i][j] += val0;
	}
      }
    }
    */

  }

  void StandardZOA::calculateElementVector(const ElInfo *elInfo, ElementVector *vec)
  {
    int nPoints = quadrature->getNumPoints();
    std::vector<double> c(nPoints, 0.0);

    int myRank = omp_get_thread_num();
    std::vector<OperatorTerm*>::iterator termIt;
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt) {
      (static_cast<ZeroOrderTerm*>((*termIt)))->getC(elInfo, nPoints, c);
    }

    for (int iq = 0; iq < nPoints; iq++) {
      c[iq] *= elInfo->getDet();

      for (int i = 0; i < nRow; i++) {
	double psi = (*(owner->getRowFESpace()->getBasisFcts()->getPhi(i)))
	  (quadrature->getLambda(iq));
	(*vec)[i] += quadrature->getWeight(iq) * c[iq] * psi;
      }
    }
  }

  FastQuadZOA::FastQuadZOA(Operator *op, Assembler *assembler, Quadrature *quad)
    : ZeroOrderAssembler(op, assembler, quad, true)
  {}

  void FastQuadZOA::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix *mat)
  {
    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();

    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	const BasisFunction *basFcts = owner->getRowFESpace()->getBasisFcts();
	psiFast = updateFastQuadrature(psiFast, basFcts, INIT_PHI);
	basFcts = owner->getColFESpace()->getBasisFcts();
	phiFast = updateFastQuadrature(phiFast, basFcts, INIT_PHI);
	firstCall = false;
      }
    }

    std::vector<double> c(nPoints, 0.0);
    std::vector<OperatorTerm*>::iterator termIt;
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt) {
      (static_cast<ZeroOrderTerm*>((*termIt)))->getC(elInfo, nPoints, c);
    }

    if (symmetric) {
      TEST_EXIT_DBG(nCol == nRow)("nCol != nRow, but symmetric assembling!\n");

      for (int iq = 0; iq < nPoints; iq++) {
	c[iq] *= elInfo->getDet();

	const double *psi = psiFast->getPhi(iq);
	const double *phi = phiFast->getPhi(iq);
	for (int i = 0; i < nRow; i++) {
	  (*mat)[i][i] += quadrature->getWeight(iq) * c[iq] * psi[i] * phi[i];
	  for (int j = i + 1; j < nCol; j++) {
	    double val = quadrature->getWeight(iq) * c[iq] * psi[i] * phi[j];
	    (*mat)[i][j] += val;
	    (*mat)[j][i] += val;
	  }
	}
      }
    } else {      /*  non symmetric assembling   */
      for (int iq = 0; iq < nPoints; iq++) {
	c[iq] *= elInfo->getDet();

	const double *psi = psiFast->getPhi(iq);
	const double *phi = phiFast->getPhi(iq);
	for (int i = 0; i < nRow; i++) {
	  for (int j = 0; j < nCol; j++) {
	    (*mat)[i][j] += quadrature->getWeight(iq) * c[iq] * psi[i] * phi[j];
	  }
	}
      }
    }
  }

  void FastQuadZOA::calculateElementMatrix(const ElInfo *rowElInfo,
					   const ElInfo *colElInfo,
					   const ElInfo *smallElInfo,
					   const ElInfo *largeElInfo,
					   ElementMatrix *mat) 
  {
    FUNCNAME("FastQuadZOA::calculateElementMatrix()");

    ERROR_EXIT("CHECK FIRST, IF IT WILL REALLY WORK!\n");

    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    const mtl::dense2D<double>& m = smallElInfo->getSubElemCoordsMat();

    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	const BasisFunction *basFcts = owner->getRowFESpace()->getBasisFcts();
	psiFast = updateFastQuadrature(psiFast, basFcts, INIT_PHI);
	basFcts = owner->getColFESpace()->getBasisFcts();
	phiFast = updateFastQuadrature(phiFast, basFcts, INIT_PHI);
	firstCall = false;
      }
    }

    std::vector<double> c(nPoints, 0.0);
    std::vector<OperatorTerm*>::iterator termIt;
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt) {
      (static_cast<ZeroOrderTerm*>((*termIt)))->getC(rowElInfo, nPoints, c);
    }

    for (int iq = 0; iq < nPoints; iq++) {
      c[iq] *= smallElInfo->getDet();
      
      const double *psi = psiFast->getPhi(iq);
      const double *phi = phiFast->getPhi(iq);

      for (int i = 0; i < nCol; i++) { 
	for (int j = 0; j < nRow; j++) { 
	  double val = quadrature->getWeight(iq) * c[iq] * psi[i];

	  double tmpval = 0.0;
	  for (int k = 0; k < nCol; k++) {
	    tmpval += m[j][k] * phi[k];
	  }
	  val *= tmpval;

	  (*mat)[j][i] += val;
	}
      }

    }    
  }

  void FastQuadZOA::calculateElementVector(const ElInfo *elInfo, ElementVector *vec)
  {
    int myRank = omp_get_thread_num();
    int nPoints = quadrature->getNumPoints();

    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	const BasisFunction *basFcts = owner->getRowFESpace()->getBasisFcts();
	psiFast = updateFastQuadrature(psiFast, basFcts, INIT_PHI);
	basFcts = owner->getColFESpace()->getBasisFcts();
	phiFast = updateFastQuadrature(phiFast, basFcts, INIT_PHI);
	firstCall = false;
      }
    }

    std::vector<double> c(nPoints, 0.0);
    std::vector<OperatorTerm*>::iterator termIt;
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt) {
      (static_cast<ZeroOrderTerm*>((*termIt)))->getC(elInfo, nPoints, c);
    }

    for (int iq = 0; iq < nPoints; iq++) {
      c[iq] *= elInfo->getDet();

      const double *psi = psiFast->getPhi(iq);
      for (int i = 0; i < nRow; i++) {
	(*vec)[i] += quadrature->getWeight(iq) * c[iq] * psi[i];
      }
    }
  }

  PrecalcZOA::PrecalcZOA(Operator *op, Assembler *assembler, Quadrature *quad) 
    : ZeroOrderAssembler(op, assembler, quad, true)
  {
  }

  void PrecalcZOA::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix *mat)
  {
    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	q00 = Q00PsiPhi::provideQ00PsiPhi(owner->getRowFESpace()->getBasisFcts(), 
					  owner->getColFESpace()->getBasisFcts(), 
					  quadrature);
	q0 = Q0Psi::provideQ0Psi(owner->getRowFESpace()->getBasisFcts(),
				 quadrature);
	firstCall = false;
      }
    }

    std::vector<double> c(1, 0.0);
    int myRank = omp_get_thread_num();
    int size = static_cast<int>(terms[myRank].size());

    for (int i = 0; i < size; i++) {
      (static_cast<ZeroOrderTerm*>((terms[myRank][i])))->getC(elInfo, 1, c);
    }

    c[0] *= elInfo->getDet();

    if (symmetric) {
      for (int i = 0; i < nRow; i++) {
	(*mat)[i][i] += c[0] * q00->getValue(i,i);
	for (int j = i + 1; j < nCol; j++) {
	  double val = c[0] * q00->getValue(i, j);
	  (*mat)[i][j] += val;
	  (*mat)[j][i] += val;
	}
      }
    } else {
      for (int i = 0; i < nRow; i++) {
	for (int j = 0; j < nCol; j++) {
	  (*mat)[i][j] += c[0] * q00->getValue(i, j);
	}
      }
    }
  }

  void PrecalcZOA::calculateElementVector(const ElInfo *elInfo, ElementVector *vec)
  {
    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	q00 = Q00PsiPhi::provideQ00PsiPhi(owner->getRowFESpace()->getBasisFcts(), 
					  owner->getColFESpace()->getBasisFcts(), 
					  quadrature);
	q0 = Q0Psi::provideQ0Psi(owner->getRowFESpace()->getBasisFcts(),
				 quadrature);	
	firstCall = false;
      }
    }

    std::vector<OperatorTerm*>::iterator termIt;

    int myRank = omp_get_thread_num();
    std::vector<double> c(1, 0.0);
    for (termIt = terms[myRank].begin(); termIt != terms[myRank].end(); ++termIt)
      (static_cast<ZeroOrderTerm*>( *termIt))->getC(elInfo, 1, c);

    c[0] *= elInfo->getDet();

    for (int i = 0; i < nRow; i++)
      (*vec)[i] += c[0] * q0->getValue(i);
  }

}