#include "RobinBC.h"
#include "Estimator.h"
#include "Assembler.h"
#include "DOFVector.h"
#include "DOFMatrix.h"
#include "OpenMP.h"
#include "SurfaceOperator.h"
#include <math.h>

namespace AMDiS {

  RobinBC::RobinBC(BoundaryType type,
		   AbstractFunction<double, WorldVector<double> > *j,
		   AbstractFunction<double, WorldVector<double> > *alpha,
		   FiniteElemSpace *rowFESpace_,
		   FiniteElemSpace *colFESpace_)
    : BoundaryCondition(type, rowFESpace_, colFESpace_), 
      neumannOperators(NULL), 
      robinOperators(NULL)
		  //matrix(dofMatrix)
  {
    int i, k, dim = rowFESpace->getMesh()->getDim();

    // create barycentric coords for each vertex of each side
    const Element *refElement = Global::getReferenceElement(dim);

    coords = new VectorOfFixVecs<DimVec<double> >*[dim+1];

    // for all element sides
    for(i=0; i < dim+1; i++) {
      coords[i] = 
	new VectorOfFixVecs<DimVec<double> >(dim,
					     dim,
					     DEFAULT_VALUE,
					     DimVec<double>(dim, 
							    DEFAULT_VALUE, 
							    0.0));
      // for each vertex of the side
      for(k=0; k < dim; k++) {
	int index = 
	  refElement->getVertexOfPosition(INDEX_OF_DIM(dim-1, dim), i, k);
	(*coords[i])[k][index] = 1.0;
      }
    }

    if(j) {
      Operator *jOp = new Operator(Operator::VECTOR_OPERATOR, rowFESpace);
      jOp->addZeroOrderTerm(new CoordsAtQP_ZOT(j));
      neumannOperators = new DimVec<SurfaceOperator*>(dim, NO_INIT);
    
      for(i=0; i < dim+1; i++) {
	(*neumannOperators)[i] = new SurfaceOperator(jOp, *coords[i]);
      }

      delete jOp;
    }

    if(alpha) {
      Operator *alphaOp = new Operator(Operator::MATRIX_OPERATOR, 
				       rowFESpace, colFESpace);
      alphaOp->addZeroOrderTerm(new CoordsAtQP_ZOT(alpha));
      robinOperators = new DimVec<SurfaceOperator*>(dim, NO_INIT);

      for(i=0; i < dim + 1; i++) {
	(*robinOperators)[i] = new SurfaceOperator(alphaOp, *coords[i]);
      }

      delete alphaOp;
    }
  }

  RobinBC::RobinBC(BoundaryType type,
		   DOFVectorBase<double> *j,
		   DOFVectorBase<double> *alpha,
		   FiniteElemSpace *rowFESpace_,
		   FiniteElemSpace *colFESpace_)
    : BoundaryCondition(type, rowFESpace_, colFESpace_), 
      neumannOperators(NULL), 
      robinOperators(NULL)
		  //matrix(dofMatrix)
  {
    int i, k, dim = rowFESpace->getMesh()->getDim();

    // create barycentric coords for each vertex of each side
    const Element *refElement = Global::getReferenceElement(dim);

    coords = new VectorOfFixVecs<DimVec<double> >*[dim+1];

    // for all element sides
    for(i=0; i < dim+1; i++) {
      coords[i] = 
	new VectorOfFixVecs<DimVec<double> >(dim, 
					     dim, 
					     DEFAULT_VALUE, 
					     DimVec<double>(dim, 
							    DEFAULT_VALUE, 
							    0.0)
					     );
      // for each vertex of the side
      for(k=0; k < dim; k++) {
	int index = 
	  refElement->getVertexOfPosition(INDEX_OF_DIM(dim-1, dim), i, k);
	(*coords[i])[k][index] = 1.0;
      }
    }

    if(j) {
      Operator *jOp = new Operator(Operator::VECTOR_OPERATOR, rowFESpace);
      jOp->addZeroOrderTerm(new VecAtQP_ZOT(j, NULL));
      neumannOperators = new DimVec<SurfaceOperator*>(dim, NO_INIT);
    
      for(i=0; i < dim+1; i++) {
	(*neumannOperators)[i] = new SurfaceOperator(jOp, *coords[i]);
      }

      delete jOp;
    }

    if(alpha) {
      Operator *alphaOp = new Operator(Operator::MATRIX_OPERATOR, 
				       rowFESpace, colFESpace);
      alphaOp->addZeroOrderTerm(new VecAtQP_ZOT(alpha, NULL));
      robinOperators = new DimVec<SurfaceOperator*>(dim, NO_INIT);

      for(i=0; i < dim + 1; i++) {
	(*robinOperators)[i] = new SurfaceOperator(alphaOp, *coords[i]);
      }

      delete alphaOp;
    }
  }

  RobinBC::RobinBC(BoundaryType type,
		   Operator* jOp, Operator* alphaOp,
		   FiniteElemSpace *rowFESpace_,
		   FiniteElemSpace *colFESpace_) 
    : BoundaryCondition(type, rowFESpace_, colFESpace_), 
      neumannOperators(NULL), 
      robinOperators(NULL)
		  //matrix(dofMatrix)
  {
    int i, k, dim = rowFESpace->getMesh()->getDim();

    // create barycentric coords for each vertex of each side
    const Element *refElement = Global::getReferenceElement(dim);

    coords = new VectorOfFixVecs<DimVec<double> >*[dim+1];

    // for all element sides
    for(i=0; i < dim+1; i++) {
      coords[i] = 
	new VectorOfFixVecs<DimVec<double> >(dim,
					     dim, 
					     DEFAULT_VALUE,
					     DimVec<double>(dim,
							    DEFAULT_VALUE, 
							    0.0));
      // for each vertex of the side
      for(k=0; k < dim; k++) {
	int index = 
	  refElement->getVertexOfPosition(INDEX_OF_DIM(dim-1, dim), i, k);
	(*coords[i])[k][index] = 1.0;
      }
    }

    neumannOperators = new DimVec<SurfaceOperator*>(dim, NO_INIT);
    robinOperators = new DimVec<SurfaceOperator*>(dim, NO_INIT);

    for(i=0; i < dim + 1; i++) {
      if(jOp)
	(*neumannOperators)[i] = new SurfaceOperator(jOp, *coords[i]);
      if(alphaOp)
	(*robinOperators)[i] = new SurfaceOperator(alphaOp, *coords[i]);
    }   
  }

  void RobinBC::fillBoundaryCondition(DOFVectorBase<double>* vector, 
				      ElInfo* elInfo,
				      const DegreeOfFreedom* dofIndices,
				      const BoundaryType* localBound,
				      int nBasFcts)
  {
    FUNCNAME("RobinBC::fillBoundaryCondition()");
    TEST_EXIT_DBG(vector->getFESpace() == rowFESpace)("invalid row fe space\n");

    int dim = elInfo->getMesh()->getDim();

    if (neumannOperators) {
      for (int i = 0; i < dim + 1; i++) {
	if (elInfo->getBoundary(i) == boundaryType) {
	  vector->assemble(1.0, elInfo, localBound, (*neumannOperators)[i]);
	}
      }
    }
  }

  void RobinBC::fillBoundaryCondition(DOFMatrix* matrix,
				      ElInfo* elInfo,
				      const DegreeOfFreedom* dofIndices,
				      const BoundaryType* localBound,
				      int nBasFcts) 
  {
    int dim = elInfo->getMesh()->getDim();

    if (robinOperators) {
      for (int i = 0; i < dim + 1; i++) {
	if (elInfo->getBoundary(i) == boundaryType) {
	  matrix->assemble(-1.0, elInfo, localBound, (*robinOperators)[i]);
	}
      }
    }
  }
  
  double RobinBC::boundResidual(ElInfo *elInfo,
				DOFMatrix *matrix,
				const DOFVectorBase<double> *dv)
  {
    FUNCNAME("RobinBC::fillBoundaryCondition()");
    TEST_EXIT(matrix->getRowFESpace() == rowFESpace)
      ("invalid row fe space\n");
    TEST_EXIT(matrix->getColFESpace() == colFESpace)
      ("invalid col fe space\n");

    int dim = elInfo->getMesh()->getDim();
    int iq;
    DimVec<double>  lambda(dim, NO_INIT);
    double n_A_grdUh, val = 0.0;
    WorldVector<double> normal;

    const DimVec<WorldVector<double> > &Lambda = elInfo->getGrdLambda();
    double det = elInfo->getDet();

    int numPoints;

    bool neumannQuad = false;

    const BasisFunction *basFcts = dv->getFESpace()->getBasisFcts();

    TEST_EXIT(basFcts == rowFESpace->getBasisFcts())("invalid basFcts\n");

    double *uhEl = new double[basFcts->getNumber()];

    dv->getLocalVector(elInfo->getElement(), uhEl);

    TEST_EXIT(neumannOperators || robinOperators)
      ("neither neumann nor robin operators set!\n");

    if(!robinOperators) 
      neumannQuad = true;
    else {
      if(neumannOperators) {
	if((*neumannOperators)[0]->getAssembler(omp_get_thread_num())->
	   getZeroOrderAssembler()->getQuadrature()->getNumPoints() > 
	   (*robinOperators)[0]->getAssembler(omp_get_thread_num())->
	   getZeroOrderAssembler()->getQuadrature()->getNumPoints()) 
	  {
	    neumannQuad = true;
	  }
      }
    }

    Quadrature *quadrature = NULL;

    int face;

    std::vector<Operator*>::iterator op;
    for(op=matrix->getOperatorsBegin(); op != matrix->getOperatorsEnd(); ++op) {
      (*op)->getAssembler(omp_get_thread_num())->initElement(elInfo);
    }

    for(face = 0; face < dim+1; face++) {

      elInfo->getNormal(face, normal);

      quadrature = 
	neumannQuad ? 
	(*neumannOperators)[face]->getAssembler(omp_get_thread_num())->
	getZeroOrderAssembler()->getQuadrature() :
	(*robinOperators)[face]->getAssembler(omp_get_thread_num())->
	getZeroOrderAssembler()->getQuadrature();

      numPoints = quadrature->getNumPoints();

      if(elInfo->getBoundary(face) == boundaryType) {

	(*neumannOperators)[face]->getAssembler(omp_get_thread_num())->getZeroOrderAssembler()->
	  initElement(elInfo);

	//const double *uhAtQp = quadrature->uhAtQp(basFcts, uhEl, NULL);
	const double *uhAtQp = dv->getVecAtQPs(elInfo,
					       quadrature,
					       NULL,
					       NULL);

	std::vector<double> f(numPoints, 0.0);

	if(robinOperators) {
	  (*robinOperators)[face]->evalZeroOrder(numPoints, 
						 uhAtQp,
						 NULL,
						 NULL,
						 &f[0],
						 1.0);
	}

	WorldVector<double> *grdUh = new WorldVector<double>[numPoints];
	WorldVector<double> *A_grdUh = new WorldVector<double>[numPoints];

	for (iq = 0; iq < numPoints; iq++) {
	  A_grdUh[iq].set(0.0);	
	  lambda = quadrature->getLambda(iq);
	  basFcts->evalGrdUh(lambda, Lambda, uhEl, &grdUh[iq]);
	}

	for(op=matrix->getOperatorsBegin(); op != matrix->getOperatorsEnd(); ++op) {
	  (*op)->weakEvalSecondOrder(numPoints, 
				     grdUh, 
				     A_grdUh);
	}


	if (neumannOperators)
	  (*neumannOperators)[face]->getC(elInfo, numPoints, f);

	for (val = iq = 0; iq < numPoints; iq++) {
	  n_A_grdUh = (normal*A_grdUh[iq]) - f[iq]; 
	  val += quadrature->getWeight(iq)*sqr(n_A_grdUh);
	}

	delete [] grdUh;
	delete [] A_grdUh;
      }
    }

    delete [] uhEl;

    return det * val;
  }

}