#include <vector>
#include "Assembler.h"
#include "SecondOrderAssembler.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*> SecondOrderAssembler::optimizedSubAssemblers; 
  std::vector<SubAssembler*> SecondOrderAssembler::standardSubAssemblers;

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

  SecondOrderAssembler* 
  SecondOrderAssembler::getSubAssembler(Operator* op,
					Assembler *assembler,
					Quadrature *quad,
					bool optimized)
  {
    int myRank = omp_get_thread_num();

    // check if a assembler is needed at all
    if (op->secondOrder[myRank].size() == 0) {
      return NULL;
    }

    SecondOrderAssembler *newAssembler;

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

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

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

    // 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<SecondOrderAssembler*>((*subAssemblers)[i]);
      }
    }

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

    // create new assembler
    //    if (!optimized) {
      newAssembler = NEW Stand2(op, assembler, quad);
//     } else {
//       if (pwConst) {
// 	newAssembler = NEW Pre2(op, assembler, quad);
//       } else {
// 	newAssembler = NEW Quad2(op, assembler, quad);
//       }
//     }

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

  Pre2::Pre2(Operator *op, Assembler *assembler, Quadrature *quad) 
    : SecondOrderAssembler(op, assembler, quad, true)
  {
    q11 = Q11PsiPhi::provideQ11PsiPhi(owner->getRowFESpace()->getBasisFcts(), 
				      owner->getColFESpace()->getBasisFcts(), 
				      quadrature);
    tmpLALt.resize(omp_get_overall_max_threads());
    for (int i = 0; i < omp_get_overall_max_threads(); i++) {
      tmpLALt[i] = NEW DimMat<double>*;
      *(tmpLALt[i]) = NEW DimMat<double>(dim, NO_INIT);
    }
  }

  Pre2::~Pre2()
  {
    for (int i = 0; i < static_cast<int>(tmpLALt.size()); i++) {
      DELETE *(tmpLALt[i]);
      DELETE tmpLALt[i];
    }
  }

  void Pre2::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix *mat)
  {
    const int **nEntries;
    const int *k, *l;
    const double *values;

    int myRank = omp_get_thread_num();
    DimMat<double> **LALt = tmpLALt[0];
    DimMat<double> &tmpMat = *LALt[0];
    tmpMat.set(0.0);

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

    tmpMat *= elInfo->getDet();
    nEntries = q11->getNumberEntries();

    if (symmetric) {
      for (int i = 0; i < nRow; i++) {
	k = q11->getKVec(i, i);
	l = q11->getLVec(i, i);
	values = q11->getValVec(i, i);
	double val = 0.0;
	for (int m = 0; m < nEntries[i][i]; m++) {
	  val += values[m] * tmpMat[k[m]][l[m]];
	}
	(*mat)[i][i] += val;

	for (int j = i + 1; j < nCol; j++) {
	  k = q11->getKVec(i, j);
	  l = q11->getLVec(i, j);
	  values = q11->getValVec(i, j);
	  val = 0.0;
	  for (int m = 0; m < nEntries[i][j]; m++) {
	    val += values[m] * tmpMat[k[m]][l[m]];
	  }
	  (*mat)[i][j] += val;
	  (*mat)[j][i] += val;
	}
      }
    } else {  /*  A not symmetric or psi != phi        */
      for (int i = 0; i < nRow; i++) {
	for (int j = 0; j < nCol; j++) {
	  k = q11->getKVec(i, j);
	  l = q11->getLVec(i, j);
	  values = q11->getValVec(i, j);
	  double val = 0.0;
	  for (int m = 0; m < nEntries[i][j]; m++) {
	    val += values[m] * tmpMat[k[m]][l[m]];
	  }
	  (*mat)[i][j] += val;
	}
      }
    }
  }

  Quad2::Quad2(Operator *op, Assembler *assembler, Quadrature *quad) 
    : SecondOrderAssembler(op, assembler, quad, true)
  {
    tmpLALt.resize(omp_get_overall_max_threads());
  }

  Quad2::~Quad2()
  {
    if (!firstCall) {
      int nPoints = quadrature->getNumPoints();
      for (int i = 0; i < static_cast<int>(tmpLALt.size()); i++) {
	for (int j = 0; j < nPoints; j++) {
	  DELETE tmpLALt[i][j];
	}
	DELETE [] tmpLALt[i];
      }
    }
  }

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

    if (firstCall) {
      tmpLALt[myRank] = NEW DimMat<double>*[nPoints];
      for (int j = 0; j < nPoints; j++) {
	tmpLALt[myRank][j] = NEW DimMat<double>(dim, NO_INIT);
      }
    
      const BasisFunction *basFcts = owner->getRowFESpace()->getBasisFcts();
      psiFast = updateFastQuadrature(psiFast, basFcts, INIT_GRD_PHI);
      basFcts = owner->getColFESpace()->getBasisFcts();
      phiFast = updateFastQuadrature(phiFast, basFcts, INIT_GRD_PHI);
      firstCall = false;
    }

    DimMat<double> **LALt = tmpLALt[myRank];
    
    for (int i = 0; i < nPoints; i++) {
      LALt[i]->set(0.0);
    }
    
    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++) {
      (static_cast<SecondOrderTerm*>(terms[myRank][i]))->getLALt(elInfo, nPoints, LALt);
    }
        
    VectorOfFixVecs< DimVec<double> > *grdPsi, *grdPhi;

    if (symmetric) {
      for (int iq = 0; iq < nPoints; iq++) {
	(*LALt[iq]) *= elInfo->getDet();
	
	grdPsi = psiFast->getGradient(iq);
	grdPhi = phiFast->getGradient(iq);

	for (int i = 0; i < nRow; i++) {
	  (*mat)[i][i] += quadrature->getWeight(iq) * 
	    xAy((*grdPsi)[i], (*LALt[iq]), (*grdPhi)[i]);

	  for (int j = i + 1; j < nCol; j++) {
	    double val = quadrature->getWeight(iq) * 
	      xAy((*grdPsi)[i], (*LALt[iq]), (*grdPhi)[j]);
	    (*mat)[i][j] += val;
	    (*mat)[j][i] += val;
	  }
	}
      }
    } else {      /*  non symmetric assembling   */
      for (int iq = 0; iq < nPoints; iq++) {
	(*LALt[iq]) *= elInfo->getDet();

	grdPsi = psiFast->getGradient(iq);
	grdPhi = phiFast->getGradient(iq);

	for (int i = 0; i < nRow; i++) {
	  for (int j = 0; j < nCol; j++) {
	    (*mat)[i][j] += quadrature->getWeight(iq) * 
	      xAy((*grdPsi)[i], (*LALt[iq]), (*grdPhi)[j]);
	  }
	}
      }
    }
  }

  Stand2::Stand2(Operator *op, Assembler *assembler, Quadrature *quad) 
    : SecondOrderAssembler(op, assembler, quad, false)
  {}

  void Stand2::calculateElementMatrix(const ElInfo *elInfo, ElementMatrix *mat)
  {
    DimVec<double> grdPsi(dim, NO_INIT);
    VectorOfFixVecs<DimVec<double> > grdPhi(dim, nCol, NO_INIT);

    const BasisFunction *psi = owner->getRowFESpace()->getBasisFcts();
    const BasisFunction *phi = owner->getColFESpace()->getBasisFcts();

    int nPoints = quadrature->getNumPoints();

    DimMat<double> **LALt = NEW DimMat<double>*[nPoints];
    for (int iq = 0; iq < nPoints; iq++) {
      LALt[iq] = NEW DimMat<double>(dim, NO_INIT);
      LALt[iq]->set(0.0);
    }

    int myRank = omp_get_thread_num();

    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++) {
      (static_cast<SecondOrderTerm*>(terms[myRank][i]))->getLALt(elInfo, nPoints, LALt);
    }

    if (symmetric) {
      for (int iq = 0; iq < nPoints; iq++) {
	(*LALt[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++) {
	  (*(psi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPsi);
	  (*mat)[i][i] += quadrature->getWeight(iq) * 
	    (grdPsi * ((*LALt[iq]) * grdPhi[i]));

	  for (int j = i + 1; j < nCol; j++) {
	    double val = quadrature->getWeight(iq) * 
	      (grdPsi * ((*LALt[iq]) * grdPhi[j]));
	    (*mat)[i][j] += val;
	    (*mat)[j][i] += val;
	  }
	}
      }
    } else {      /*  non symmetric assembling   */
      for (int iq = 0; iq < nPoints; iq++) {
	(*LALt[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++) {
	  (*(psi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPsi);
	  for (int j = 0; j < nCol; j++) {
	    (*mat)[i][j] += quadrature->getWeight(iq) *
	      (grdPsi * ((*LALt[iq]) * grdPhi[j]));
	  }
	}
      }
    }

    for (int iq = 0; iq < nPoints; iq++) 
      DELETE LALt[iq];

    DELETE [] LALt;
  }

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

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

    bool p = false;
    /*
    if (smallElInfo->getLevel() == largeElInfo->getLevel() + 1) {
      p = true;
      std::cout << "--------------------------------------" << std::endl;
      std::cout << "i-th child = " << smallElInfo->getIChild() << std::endl;
      std::cout << "Mesh row = " << owner->getRowFESpace()->getMesh() << std::endl;
      std::cout << "Mesh col = " << owner->getColFESpace()->getMesh() << std::endl;
      std::cout << "Mesh main = " << rowElInfo->getMesh() << std::endl;
      std::cout << "Mesh aux = " << colElInfo->getMesh() << std::endl;
    }
    */

    DimVec<double> grdPsi(dim, NO_INIT);
    VectorOfFixVecs<DimVec<double> > grdPhi(dim, nCol, NO_INIT);

    const BasisFunction *psi = owner->getRowFESpace()->getBasisFcts();
    const BasisFunction *phi = owner->getColFESpace()->getBasisFcts();
    DimMat<double> *m = smallElInfo->getSubElemCoordsMat();

    TEST_EXIT(m)("No subElemCoordsMat!\n");

    int nPoints = quadrature->getNumPoints();

    DimMat<double> **LALt = NEW DimMat<double>*[nPoints];
    for (int iq = 0; iq < nPoints; iq++) {
      LALt[iq] = NEW DimMat<double>(dim, NO_INIT);
      LALt[iq]->set(0.0);
    }

    int myRank = omp_get_thread_num();

    if (p) {
      std::cout << "Terms = " << terms[myRank].size() << std::endl;
    }

    for (int i = 0; i < static_cast<int>(terms[myRank].size()); i++) {
      (static_cast<SecondOrderTerm*>(terms[myRank][i]))
	->getLALt(smallElInfo, nPoints, LALt);
    }

    if (p) {
      for (int i = 0; i < nPoints; i++) {
	std::cout << "LALt for iq = " << i << std::endl;
	LALt[i]->print();
      }
    }

    DimVec<double> val1(dim, DEFAULT_VALUE, 0.0);

    if (p)
      std::cout << "Multiply LALt with det = " << smallElInfo->getDet() << std::endl;

    for (int iq = 0; iq < nPoints; iq++) {
      (*LALt[iq]) *= smallElInfo->getDet();
    }
     
    /*
    for (int i = 0; i < nRow; i++) {
      for (int j = 0; j < nCol; j++) {	
	for (int iq = 0; iq < nPoints; iq++) {

	  val1.set(0.0);

	  for (int k = 0; k < nCol; k++) {
	    (*(phi->getGrdPhi(k)))(quadrature->getLambda(iq), grdPsi);
	    grdPsi *= (*m)[j][k];
	    val1 += grdPsi;
	  }

	  (*(psi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPsi);

	  (*mat)[i][j] += quadrature->getWeight(iq) *
	      (grdPsi * ((*LALt[iq]) * val1));
	}
      }
      }*/

      for (int iq = 0; iq < nPoints; iq++) {
	for (int i = 0; i < nCol; i++) {
	  (*(phi->getGrdPhi(i)))(quadrature->getLambda(iq), grdPhi[i]);
	  if (p) {
	    std::cout << "grdPhi[" << i << "]" << std::endl;
	    grdPhi[i].print();
	  }
	}

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

      if (p) {
	mat->print();
	WAIT_REALLY;
      }


    for (int iq = 0; iq < nPoints; iq++) 
      DELETE LALt[iq];

    DELETE [] LALt;
  }

}