#include <vector>
#include <boost/numeric/mtl/mtl.hpp>
#include "Assembler.h"
#include "FirstOrderAssembler.h"
#include "Operator.h"
#include "QPsiPhi.h"
#include "FiniteElemSpace.h"
#include "Quadrature.h"
#include "DOFVector.h"
#include "OpenMP.h"

namespace AMDiS {

  std::vector<SubAssembler*> FirstOrderAssembler::optimizedSubAssemblersGrdPhi;
  std::vector<SubAssembler*> FirstOrderAssembler::optimizedSubAssemblersGrdPsi;

  std::vector<SubAssembler*> FirstOrderAssembler::standardSubAssemblersGrdPhi;
  std::vector<SubAssembler*> FirstOrderAssembler::standardSubAssemblersGrdPsi;


  FirstOrderAssembler::FirstOrderAssembler(Operator *op,
					   Assembler *assembler,
					   Quadrature *quad,
					   bool optimized,
					   FirstOrderType type)
    : SubAssembler(op, assembler, quad, 1, optimized, type)
  {
    VectorOfFixVecs<DimVec<double> > 
      Lb(assembler->getRowFESpace()->getMesh()->getDim(), 1, NO_INIT);
    tmpLb.resize(omp_get_overall_max_threads(), Lb);
  }

  FirstOrderAssembler* 
  FirstOrderAssembler::getSubAssembler(Operator* op,
				       Assembler *assembler,
				       Quadrature *quad,
				       FirstOrderType type,
				       bool optimized)
  {
    std::vector<SubAssembler*> *subAssemblers =
      optimized ?
      (type == GRD_PSI ? 
       &optimizedSubAssemblersGrdPsi : 
       &optimizedSubAssemblersGrdPhi) :
      (type == GRD_PSI ? 
       &standardSubAssemblersGrdPsi :
       &standardSubAssemblersGrdPhi);
    
    int myRank = omp_get_thread_num();
    std::vector<OperatorTerm*> opTerms = 
      (type == GRD_PSI) ? op->firstOrderGrdPsi[myRank] : op->firstOrderGrdPhi[myRank];

    // check if a assembler is needed at all
    if (opTerms.size() == 0)
      return NULL;

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

    FirstOrderAssembler *newAssembler;

    // check if a new assembler is needed
    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<FirstOrderAssembler*>((*subAssemblers)[i]);
    }

    // check if all terms are pw_const
    bool pwConst = true;
    for (int i = 0; i < static_cast<int>(opTerms.size()); i++) {
      if (!(opTerms[i])->isPWConst()) {
	pwConst = false;
	break;
      }
    }  

    // create new assembler
    if (!optimized) {
      newAssembler = 
	(type == GRD_PSI) ?
	dynamic_cast<FirstOrderAssembler*>(new Stand10(op, assembler, quad)) :
	dynamic_cast<FirstOrderAssembler*>(new Stand01(op, assembler, quad));    
    } else {
       if (pwConst) {
 	newAssembler = 
 	  (type == GRD_PSI) ?
 	  dynamic_cast<FirstOrderAssembler*>(new Pre10(op, assembler, quad)) :
 	  dynamic_cast<FirstOrderAssembler*>(new Pre01(op, assembler, quad));
       } else {
 	newAssembler = 
 	  (type == GRD_PSI) ?
 	  dynamic_cast<FirstOrderAssembler*>(new Quad10(op, assembler, quad)) :
 	  dynamic_cast<FirstOrderAssembler*>(new Quad01(op, assembler, quad));
       }
    }

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

  Stand10::Stand10(Operator *op, Assembler *assembler, Quadrature *quad) 
    : FirstOrderAssembler(op, assembler, quad, false, GRD_PSI)
  {
    psi = owner->getRowFESpace()->getBasisFcts();
    phi = owner->getColFESpace()->getBasisFcts();
  }

  void Stand10::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix& mat)
  {
    DimVec<double> grdPsi(dim, DEFAULT_VALUE, 0.0);

    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb.resize(nPoints);
    std::vector<double> phival(nCol);

    for (int iq = 0; iq < nPoints; iq++)
      Lb[iq].set(0.0);    
    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++)
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, nPoints, Lb);    
  
    for (int iq = 0; iq < nPoints; iq++) {
      Lb[iq] *= elInfo->getDet();

      for (int i = 0; i < nCol; i++)
	phival[i] = (*(phi->getPhi(i)))(quadrature->getLambda(iq));      

      for (int i = 0; i < nRow; i++) {
	(*(psi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPsi);
	for (int j = 0; j < nCol; j++)
	  mat[i][j] += quadrature->getWeight(iq) * phival[j] * (Lb[iq] * grdPsi);
      }
    }
  }

  void Stand10::calculateElementVector(const ElInfo *elInfo, ElementVector& vec)
  {
    DimVec<double> grdPsi(dim, DEFAULT_VALUE, 0.0);
    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb.resize(nPoints);

    for (int iq = 0; iq < nPoints; iq++)
      Lb[iq].set(0.0);    
    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++)
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, nPoints, Lb);    
  
    for (int iq = 0; iq < nPoints; iq++) {
      Lb[iq] *= elInfo->getDet();

      for (int i = 0; i < nRow; i++) {
	(*(psi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPsi);
	vec[i] += quadrature->getWeight(iq) * (Lb[iq] * grdPsi);
      }
    }
  }


  Quad10::Quad10(Operator *op, Assembler *assembler, Quadrature *quad) 
    : FirstOrderAssembler(op, assembler, quad, true, GRD_PSI)
  {}


  void Quad10::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix& mat)
  {
    VectorOfFixVecs<DimVec<double> > *grdPsi;

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

    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb.resize(nPoints);

    for (int iq = 0; iq < nPoints; iq++)
      Lb[iq].set(0.0);

    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++)
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, nPoints, Lb);
  
    for (int iq = 0; iq < nPoints; iq++) {
      Lb[iq] *= elInfo->getDet();

      const double *phi = phiFast->getPhi(iq);
      grdPsi = psiFast->getGradient(iq);
      double factor = quadrature->getWeight(iq);

      for (int i = 0; i < nRow; i++) {
	for (int j = 0; j < nCol; j++)
	  mat[i][j] += factor * (Lb[iq] * (*grdPsi)[i]) * phi[j];
      }
    }
  }

  void Quad10::calculateElementVector(const ElInfo *elInfo, ElementVector& vec)
  {
    VectorOfFixVecs<DimVec<double> > *grdPsi;

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

    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb.resize(nPoints);

    for (int iq = 0; iq < nPoints; iq++)
      Lb[iq].set(0.0);    
    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++)
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, nPoints, Lb);
  
    for (int iq = 0; iq < nPoints; iq++) {
      Lb[iq] *= elInfo->getDet();
      grdPsi = psiFast->getGradient(iq);

      for (int i = 0; i < nRow; i++)
	vec[i] += quadrature->getWeight(iq) * (Lb[iq] * (*grdPsi)[i]);
    }
  }

  Pre10::Pre10(Operator *op, Assembler *assembler, Quadrature *quad) 
    : FirstOrderAssembler(op, assembler, quad, true, GRD_PSI)
  {
  }


  void Pre10::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix& mat)
  {
    const int *k;
    const double *values;

    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {	
	q10 = Q10PsiPhi::provideQ10PsiPhi(owner->getRowFESpace()->getBasisFcts(), 
					  owner->getColFESpace()->getBasisFcts(), 
					  quadrature);
	q1 = Q1Psi::provideQ1Psi(owner->getRowFESpace()->getBasisFcts(),
				 quadrature);
	firstCall = false;
      }
    }

    const int **nEntries = q10->getNumberEntries();
    int myRank = omp_get_thread_num();
    // Do not need do resize Lb, because it's size is always at least one.
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb[0].set(0.0);

    for (int i = 0; i < static_cast<int>( terms[myRank].size()); i++) {
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, 1, Lb);
    }

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

    for (int i = 0; i < nRow; i++) {
      for (int j = 0; j < nCol; j++) {
	k = q10->getKVec(i, j);
	values = q10->getValVec(i, j);
	double val = 0.0;
	for (int m = 0; m < nEntries[i][j]; m++)
	  val += values[m] * Lb[0][k[m]];	
	mat[i][j] += val;
      }
    }
  }

  Stand01::Stand01(Operator *op, Assembler *assembler, Quadrature *quad) 
    : FirstOrderAssembler(op, assembler, quad, false, GRD_PHI)
  {
    VectorOfFixVecs<DimVec<double> > 
      grdPhi(assembler->getRowFESpace()->getMesh()->getDim(), nCol, NO_INIT);
    tmpGrdPhi.resize(omp_get_overall_max_threads(), grdPhi);

    psi = owner->getRowFESpace()->getBasisFcts();
    phi = owner->getColFESpace()->getBasisFcts();
  }

  void Stand01::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix& mat)
  {
    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    VectorOfFixVecs<DimVec<double> >& Lb = tmpLb[myRank];
    VectorOfFixVecs<DimVec<double> >& grdPhi = tmpGrdPhi[myRank];
    Lb.resize(nPoints);

    for (int iq = 0; iq < nPoints; iq++)
      Lb[iq].set(0.0);

    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++)
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, nPoints, Lb);
  
    for (int iq = 0; iq < nPoints; iq++) {
      Lb[iq] *= elInfo->getDet();

      for (int i = 0; i < nCol; i++)
	(*(phi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPhi[i]);

      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) * psival * (Lb[iq] * grdPhi[j]);
      }
    } 
  }

  Quad01::Quad01(Operator *op, Assembler *assembler, Quadrature *quad) 
    : FirstOrderAssembler(op, assembler, quad, true, GRD_PHI)
  {}

  void Quad01::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix& mat)
  {
    VectorOfFixVecs<DimVec<double> > *grdPhi;

    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_GRD_PHI);
	firstCall = false;
      }
    }

    int nPoints = quadrature->getNumPoints();
    int myRank = omp_get_thread_num();
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb.resize(nPoints);

    for (int iq = 0; iq < nPoints; iq++)
      Lb[iq].set(0.0);

    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++)
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, nPoints, Lb);
  
    for (int iq = 0; iq < nPoints; iq++) {
      Lb[iq] *= elInfo->getDet();

      const double *psi = psiFast->getPhi(iq);
      grdPhi = phiFast->getGradient(iq);

      for (int i = 0; i < nRow; i++) {
	double factor = quadrature->getWeight(iq) * psi[i];
	for (int j = 0; j < nCol; j++)
	  mat[i][j] += factor * (Lb[iq] * (*grdPhi)[j]);
      }
    }
  }

  Pre01::Pre01(Operator *op, Assembler *assembler, Quadrature *quad) 
    : FirstOrderAssembler(op, assembler, quad, true, GRD_PHI)
  {
  }

  void Pre01::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix& mat)
  {
    const int *l;
    const double *values;

    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	q01 = Q01PsiPhi::provideQ01PsiPhi(owner->getRowFESpace()->getBasisFcts(), 
					  owner->getColFESpace()->getBasisFcts(), 
					  quadrature);
	q1 = Q1Psi::provideQ1Psi(owner->getRowFESpace()->getBasisFcts(),
				 quadrature);
	firstCall = false;
      }
    }

    const int **nEntries = q01->getNumberEntries();
    int myRank = omp_get_thread_num();
    // Do not need to resize Lb, because it's size is always at least one!
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb[0].set(0.0);

    for (int i = 0; i < static_cast<int>( terms[myRank].size()); i++) {
      (static_cast<FirstOrderTerm*>((terms[myRank][i])))->getLb(elInfo, 1, Lb);
    }

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

    for (int i = 0; i < nRow; i++) {
      for (int j = 0; j < nCol; j++) {
	l = q01->getLVec(i, j);
	values = q01->getValVec(i, j);
	double val = 0.0;
	for (int m = 0; m < nEntries[i][j]; m++) {
	  val += values[m] * Lb[0][l[m]];
	}
	mat[i][j] += val;
      }
    }
  }

  void Pre10::calculateElementVector(const ElInfo *elInfo, ElementVector& vec)
  {
    const int *k;
    const double *values;

    if (firstCall) {
#ifdef _OPENMP
#pragma omp critical
#endif 
      {
	q10 = Q10PsiPhi::provideQ10PsiPhi(owner->getRowFESpace()->getBasisFcts(), 
					  owner->getColFESpace()->getBasisFcts(), 
					  quadrature);
	q1 = Q1Psi::provideQ1Psi(owner->getRowFESpace()->getBasisFcts(),
				 quadrature);
	firstCall = false;
      }
    }

    const int *nEntries = q1->getNumberEntries();
    int myRank = omp_get_thread_num();
    // Do not need to resize Lb, because it's size is always at least one!
    VectorOfFixVecs<DimVec<double> > &Lb = tmpLb[myRank];
    Lb[0].set(0.0);

    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++) {
      (static_cast<FirstOrderTerm*>(terms[myRank][i]))->getLb(elInfo, 1, Lb);
    }

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

    for (int i = 0; i < nRow; i++) {
      k = q1->getKVec(i);
      values = q1->getValVec(i);
      double val = 0.0;
      for (int m = 0; m < nEntries[i]; m++) {
	val += values[m] * Lb[0][k[m]];
      }
      vec[i] += val;
    }
  }

}