#include "Assembler.h"
#include "Operator.h"
#include "Element.h"
#include "QPsiPhi.h"
#include "ElementMatrix.h"
#include "ElementVector.h"
#include "DOFVector.h"
#include "OpenMP.h"
#include <vector>
#include <algorithm>

namespace AMDiS {

  Assembler::Assembler(Operator  *op,
		       const FiniteElemSpace *rowFESpace_,
		       const FiniteElemSpace *colFESpace_) 
    : operat(op),
      rowFESpace(rowFESpace_),
      colFESpace(colFESpace_ ? colFESpace_ : rowFESpace_),
      nRow(rowFESpace->getBasisFcts()->getNumber()),
      nCol(colFESpace->getBasisFcts()->getNumber()),
      remember(true),
      rememberElMat(false),
      rememberElVec(false),
      elementMatrix(NULL),
      elementVector(NULL),
      lastMatEl(NULL),
      lastVecEl(NULL),
      lastTraverseId(-1)
  
  {}

  void Assembler::calculateElementMatrix(const ElInfo *elInfo, 
					 ElementMatrix *userMat,
					 double factor)
  {
    FUNCNAME("Assembler::calculateElementMatrix()");

    if (remember && ((factor != 1.0) || (operat->uhOld))) {
      rememberElMat = true;
    }
  
    if (rememberElMat && !elementMatrix)
      elementMatrix = NEW ElementMatrix(nRow, nCol);

    Element *el = elInfo->getElement();

    
    checkForNewTraverse();

    checkQuadratures();
    
    if ((el != lastMatEl && el != lastVecEl) || !operat->isOptimized()) {
      initElement(elInfo);
    }

    if (el != lastMatEl || !operat->isOptimized()) {
      if (rememberElMat) {
	elementMatrix->set(0.0);
      }
      lastMatEl = el;
    } else {
      if (rememberElMat) {
	axpy(factor, *elementMatrix, *userMat);
	return;
      }
    }
  
    ElementMatrix *mat = rememberElMat ? elementMatrix : userMat;

    if (secondOrderAssembler)
      secondOrderAssembler->calculateElementMatrix(elInfo, mat);
    if (firstOrderAssemblerGrdPsi)
      firstOrderAssemblerGrdPsi->calculateElementMatrix(elInfo, mat);
    if (firstOrderAssemblerGrdPhi)
      firstOrderAssemblerGrdPhi->calculateElementMatrix(elInfo, mat);
    if (zeroOrderAssembler)
      zeroOrderAssembler->calculateElementMatrix(elInfo, mat);

    if (rememberElMat && userMat) {
      axpy(factor, *elementMatrix, *userMat);
    }      
  }

  void Assembler::calculateElementMatrix(const ElInfo *rowElInfo,
					 const ElInfo *colElInfo,
					 const ElInfo *smallElInfo,
					 const ElInfo *largeElInfo,
					 ElementMatrix *userMat,
					 double factor)
  {
    FUNCNAME("Assembler::calculateElementMatrix()");

    if (remember && ((factor != 1.0) || (operat->uhOld))) {
      rememberElMat = true;
    }
  
    if (rememberElMat && !elementMatrix)
      elementMatrix = NEW ElementMatrix(nRow, nCol);

    Element *el = rowElInfo->getElement();

    
    //    checkForNewTraverse();
    lastVecEl = lastMatEl = NULL;

    checkQuadratures();
    
    if ((el != lastMatEl && el != lastVecEl) || !operat->isOptimized()) {
      initElement(rowElInfo);
    }

    if (el != lastMatEl || !operat->isOptimized()) {
      if (rememberElMat) {
	elementMatrix->set(0.0);
      }
      lastMatEl = el;
    } else {
      if (rememberElMat) {
	axpy(factor, *elementMatrix, *userMat);
	return;
      }
    }
  
    ElementMatrix *mat = rememberElMat ? elementMatrix : userMat;

    if (secondOrderAssembler) {
      secondOrderAssembler->calculateElementMatrix(rowElInfo, colElInfo, 
						   smallElInfo, largeElInfo, mat);
    }

    if (firstOrderAssemblerGrdPsi) {
      firstOrderAssemblerGrdPsi->calculateElementMatrix(rowElInfo, colElInfo, 
							smallElInfo, largeElInfo, mat);
    }

    if (firstOrderAssemblerGrdPhi) {
      firstOrderAssemblerGrdPhi->calculateElementMatrix(rowElInfo, colElInfo, 
							smallElInfo, largeElInfo, mat);
    }

    if (zeroOrderAssembler) {
      zeroOrderAssembler->calculateElementMatrix(rowElInfo, colElInfo, 
						 smallElInfo, largeElInfo, mat);
    }
 
    if (rememberElMat && userMat) {
      axpy(factor, *elementMatrix, *userMat);
    }      
  }

  void Assembler::calculateElementVector(const ElInfo *elInfo, 
					 ElementVector *userVec,
					 double factor)
  {
    FUNCNAME("Assembler::calculateElementVector()");

    if (remember && factor != 1.0) {
      rememberElVec = true;
    }

    if (rememberElVec && !elementVector) {
      elementVector = NEW ElementVector(nRow);
    }

    Element *el = elInfo->getElement();
    checkForNewTraverse();
    checkQuadratures();

    if ((el != lastMatEl && el != lastVecEl) || !operat->isOptimized()) {
      initElement(elInfo);
    }
    
    if (el != lastVecEl || !operat->isOptimized()) {
      if (rememberElVec) {
	elementVector->set(0.0);
      }
      lastVecEl = el;
    } else {
      if (rememberElVec) {
	axpy(factor, *elementVector, *userVec);
	return;
      }
    }
    ElementVector *vec = rememberElVec ? elementVector : userVec;
    if (operat->uhOld && remember) {
      matVecAssemble(elInfo, vec);
      if (rememberElVec) {
	axpy(factor, *elementVector, *userVec);
      }
      return;
    } 
    if (firstOrderAssemblerGrdPsi) {
      firstOrderAssemblerGrdPsi->calculateElementVector(elInfo, vec);
    }
    if (zeroOrderAssembler) {
      zeroOrderAssembler->calculateElementVector(elInfo, vec);
    }
    if (rememberElVec) {
      axpy(factor, *elementVector, *userVec);
    }      
  }

  void Assembler::matVecAssemble(const ElInfo *elInfo, ElementVector *vec)
  {
    FUNCNAME("Assembler::matVecAssemble()");

    Element *el = elInfo->getElement(); 
    const BasisFunction *basFcts = rowFESpace->getBasisFcts();
    int nBasFcts = basFcts->getNumber();
    double *uhOldLoc = new double[nBasFcts];

    operat->uhOld->getLocalVector(el, uhOldLoc);
    
    if (el != lastMatEl) {
      calculateElementMatrix(elInfo, NULL);
    }
    
    for (int i = 0; i < nBasFcts; i++) {
      double val = 0.0;
      for (int j = 0; j < nBasFcts; j++) {
	val += (*elementMatrix)[i][j] * uhOldLoc[j];
      }
      (*vec)[i] += val;
    }
    

    delete [] uhOldLoc;       
  }

  void Assembler::initElement(const ElInfo *elInfo, Quadrature *quad)
  {
    checkQuadratures();

    if (secondOrderAssembler) 
      secondOrderAssembler->initElement(elInfo, quad);
    if (firstOrderAssemblerGrdPsi)
      firstOrderAssemblerGrdPsi->initElement(elInfo, quad);
    if (firstOrderAssemblerGrdPhi)
      firstOrderAssemblerGrdPhi->initElement(elInfo, quad);
    if (zeroOrderAssembler)
      zeroOrderAssembler->initElement(elInfo, quad);
  }

  OptimizedAssembler::OptimizedAssembler(Operator  *op,
					 Quadrature *quad2,
					 Quadrature *quad1GrdPsi,
					 Quadrature *quad1GrdPhi,
					 Quadrature *quad0,
					 const FiniteElemSpace *rowFESpace_,
					 const FiniteElemSpace *colFESpace_) 
    : Assembler(op, rowFESpace_, colFESpace_)
  {
    bool opt = (rowFESpace_ == colFESpace_);

    // create sub assemblers
    secondOrderAssembler = 
      SecondOrderAssembler::getSubAssembler(op, this, quad2, opt);
    firstOrderAssemblerGrdPsi = 
      FirstOrderAssembler::getSubAssembler(op, this, quad1GrdPsi, GRD_PSI, opt);
    firstOrderAssemblerGrdPhi = 
      FirstOrderAssembler::getSubAssembler(op, this, quad1GrdPhi, GRD_PHI, opt);
    zeroOrderAssembler = 
      ZeroOrderAssembler::getSubAssembler(op, this, quad0, opt);
  }

  StandardAssembler::StandardAssembler(Operator *op,
				       Quadrature *quad2,
				       Quadrature *quad1GrdPsi,
				       Quadrature *quad1GrdPhi,
				       Quadrature *quad0,
				       const FiniteElemSpace *rowFESpace_,
				       const FiniteElemSpace *colFESpace_) 
    : Assembler(op, rowFESpace_, colFESpace_)
  {
    remember = false;

    // create sub assemblers
    secondOrderAssembler = 
      SecondOrderAssembler::getSubAssembler(op, this, quad2, false);
    firstOrderAssemblerGrdPsi = 
      FirstOrderAssembler::getSubAssembler(op, this, quad1GrdPsi, GRD_PSI, false);
    firstOrderAssemblerGrdPhi = 
      FirstOrderAssembler::getSubAssembler(op, this, quad1GrdPhi, GRD_PHI, false);
    zeroOrderAssembler = 
      ZeroOrderAssembler::getSubAssembler(op, this, quad0, false);
  }

  ElementMatrix *Assembler::initElementMatrix(ElementMatrix *elMat, 
					      const ElInfo *rowElInfo,
					      const ElInfo *colElInfo)
  {
    if (!elMat) {
      elMat = NEW ElementMatrix(nRow, nCol);
    }

    elMat->set(0.0);
         
    rowFESpace->getBasisFcts()->getLocalIndicesVec(rowElInfo->getElement(),
						   rowFESpace->getAdmin(),
						   &(elMat->rowIndices));

    if (rowFESpace == colFESpace) {
      elMat->colIndices = elMat->rowIndices;
    } else {
      if (colElInfo) {
	colFESpace->getBasisFcts()->getLocalIndicesVec(colElInfo->getElement(),
						       colFESpace->getAdmin(),
						       &(elMat->colIndices));
      } else {
	// If there is no colElInfo pointer, use rowElInfo the get the indices.
	colFESpace->getBasisFcts()->getLocalIndicesVec(rowElInfo->getElement(),
						       colFESpace->getAdmin(),
						       &(elMat->colIndices));
      }
    }
    
    return elMat;
  }

  ElementVector *Assembler::initElementVector(ElementVector *elVec, 
					      const ElInfo *elInfo)
  {
    if (!elVec) {
      elVec = NEW ElementVector(nRow);
    }

    elVec->set(0.0);
  
    DOFAdmin *rowAdmin = rowFESpace->getAdmin();
    Element *element = elInfo->getElement();

    rowFESpace->getBasisFcts()->getLocalIndicesVec(element, rowAdmin, &(elVec->dofIndices));

    return elVec;
  }

  void Assembler::checkQuadratures()
  { 
    if (secondOrderAssembler) {
      // create quadrature
      if (!secondOrderAssembler->getQuadrature()) {
	int dim = rowFESpace->getMesh()->getDim();
	int degree = operat->getQuadratureDegree(2);
	Quadrature *quadrature = Quadrature::provideQuadrature(dim, degree);
	secondOrderAssembler->setQuadrature(quadrature);
      }
    }
    if (firstOrderAssemblerGrdPsi) {
      // create quadrature
      if (!firstOrderAssemblerGrdPsi->getQuadrature()) {
	int dim = rowFESpace->getMesh()->getDim();
	int degree = operat->getQuadratureDegree(1, GRD_PSI);
	Quadrature *quadrature = Quadrature::provideQuadrature(dim, degree);
	firstOrderAssemblerGrdPsi->setQuadrature(quadrature);
      }
    }
    if (firstOrderAssemblerGrdPhi) {
      // create quadrature
      if (!firstOrderAssemblerGrdPhi->getQuadrature()) {
	int dim = rowFESpace->getMesh()->getDim();
	int degree = operat->getQuadratureDegree(1, GRD_PHI);
	Quadrature *quadrature = Quadrature::provideQuadrature(dim, degree);
	firstOrderAssemblerGrdPhi->setQuadrature(quadrature);
      }
    }
    if (zeroOrderAssembler) {
      // create quadrature
      if (!zeroOrderAssembler->getQuadrature()) {
	int dim = rowFESpace->getMesh()->getDim();
	int degree = operat->getQuadratureDegree(0);
	Quadrature *quadrature = Quadrature::provideQuadrature(dim, degree);
	zeroOrderAssembler->setQuadrature(quadrature);
      }
    }
  }

}