#include "LapAssembler.hpp"

void LapAssembler::assemble_elem_contribution(const MeshType& mesh,
                                              GlobalMatrixType& matrix, 
                                              GlobalVectorType& gRhs, 
				              FEManagerType& FEManager,
					      const std::pair<int,int>& elemRange)
{
  const FEManagerType::ShapeMatrixType& shapeMatrix           = FEManager.getElemShapeFunction();
  const FEManagerType::ShapeDerivMatrixType& shapeDerivMatrix = FEManager.getElemShapeFunctionDeriv();
  const FEManagerType::ElemJacobianVectorType& jacXweightVec  = FEManager.getElemJacobianWeightVec();
  const FEManagerType::ElemQpointVectorType& qpXVec           = FEManager.getElemQpointX();
  
  const FEManagerType::ShapeMatrixType& boundaryShapeMatrix           = FEManager.getBoundaryFaceElemShapeFunction();
  const FEManagerType::ShapeDerivMatrixType& boundaryShapeDerivMatrix = FEManager.getBoundaryFaceElemShapeFunctionDeriv();
  const FEManagerType::NormalVectorType& normalVecBoundary            = FEManager.getBoundaryFaceNormalVec();
  const FEManagerType::FaceJacobianVectorType& jacXweightVecBoundary  = FEManager.getBoundaryFaceJacobianWeightVec();
  const FEManagerType::FaceQpointVectorType& qpXVecBoundary           = FEManager.getBoundaryFaceQpointX();
  
  const FEManagerType::ShapeMatrixType& bShapeMatrix             = FEManager.getInternalFaceElemShapeFunction();
  const FEManagerType::ShapeMatrixType& fShapeMatrix             = FEManager.getInternalFaceNeighborShapeFunction();
  const FEManagerType::ShapeDerivMatrixType& bShapeDerivMatrix   = FEManager.getInternalFaceElemShapeFunctionDeriv();     
  const FEManagerType::ShapeDerivMatrixType& fShapeDerivMatrix   = FEManager.getInternalFaceNeighborShapeFunctionDeriv(); 
  const FEManagerType::NormalVectorType& normalVec               = FEManager.getInternalFaceNormalVec();
  const FEManagerType::FaceJacobianVectorType& jacXweightVecSide = FEManager.getInternalFaceJacobianWeightVec();
  const FEManagerType::FaceQpointVectorType& qpXVecSide          = FEManager.getInternalFaceQpointX();

  LocalMatrixType bbMatrix, bfMatrix;
  LocalVectorType fShapeDerivNormal, bShapeDerivNormal;
  LocalVectorType rhs;

  for (SpaMesh::ElementIterator elemIt(mesh, elemRange); elemIt; ++elemIt) 
  {
    int elem = elemIt.value();
    FEManager.initVolumeFE(mesh, elemIt, InputData::p_order, InputData::q_order);
    bbMatrix.setZero();
    rhs.setZero();
    for (unsigned int j = 0; j < shapeMatrix.cols(); j++)
    {
      double forcingTerm = InputData::ExactSolution::forcingTerm(qpXVec.row(j));
      for (unsigned int d = 0; d < InputData::dim; d++)
        bbMatrix.noalias() += jacXweightVec[j] * (shapeDerivMatrix(d).col(j) * (shapeDerivMatrix(d).col(j)).transpose());
      rhs.noalias() += shapeMatrix.col(j) * jacXweightVec[j] * forcingTerm;
    }
    for (SpaMesh::NeighborIterator neighborIt(mesh,elemIt); neighborIt; ++neighborIt)
    {     
      int neighbor = neighborIt.value();
      if (elem == neighbor)
      {
        FEManager.initBoundaryFaceFE(mesh, elemIt, neighborIt.index(), InputData::p_order, InputData::q_order);
	double sideVol = jacXweightVecBoundary.sum();
	double elemSize = FEManager.computeFESize(mesh, elemIt);
	const double h = (elemSize/sideVol) * (1./pow(InputData::p_order,2.));
        for (unsigned int j = 0; j < boundaryShapeMatrix.cols(); j++)
        {
          const double bc = InputData::ExactSolution::exact_solution(qpXVecBoundary.row(j));
          bShapeDerivNormal.setZero();
          for (unsigned int d = 0; d < InputData::dim; d++)
            bShapeDerivNormal.noalias() += normalVecBoundary(j)(d)*boundaryShapeDerivMatrix(d).col(j);

          bbMatrix.noalias() -= jacXweightVecBoundary[j] * (boundaryShapeMatrix.col(j) * bShapeDerivNormal.transpose());
          bbMatrix.noalias() -= jacXweightVecBoundary[j] * (bShapeDerivNormal * (boundaryShapeMatrix.col(j)).transpose());
          bbMatrix.noalias() += jacXweightVecBoundary[j] * (InputData::penalty/h) * (boundaryShapeMatrix.col(j) * (boundaryShapeMatrix.col(j)).transpose());
          rhs.noalias() -= (jacXweightVecBoundary[j] * bc) * bShapeDerivNormal;
          rhs.noalias() += (jacXweightVecBoundary[j] * InputData::penalty/h * bc) * boundaryShapeMatrix.col(j);
        }
      }
      else
      {
        FEManager.initInternalFaceFE(mesh, elemIt, neighborIt, InputData::p_order, InputData::q_order);
        const double elemSize = FEManager.computeFESize(mesh, elemIt);
        const double neighborSize = FEManager.computeFESize(mesh, neighborIt);
        const double sideVol = jacXweightVecSide.sum();
        const double h = (0.5*(elemSize+neighborSize)/(sideVol)) * (1./pow(InputData::p_order,2.));
        bfMatrix.setZero();
        for (unsigned int j = 0; j < bShapeMatrix.cols(); j++)
        {
          fShapeDerivNormal.setZero();
          bShapeDerivNormal.setZero();
          for (unsigned int d = 0; d < InputData::dim; d++)
          {
            fShapeDerivNormal.noalias() += fShapeDerivMatrix(d).col(j) * normalVec(j)(d);
            bShapeDerivNormal.noalias() += bShapeDerivMatrix(d).col(j) * normalVec(j)(d);
          }
          bbMatrix.noalias() -= 0.5 * jacXweightVecSide[j] * (bShapeMatrix.col(j) * bShapeDerivNormal.transpose());
          bbMatrix.noalias() -= 0.5 * jacXweightVecSide[j] * (bShapeDerivNormal *  (bShapeMatrix.col(j)).transpose());
          bfMatrix.noalias() -= 0.5 * jacXweightVecSide[j] * (bShapeMatrix.col(j) * fShapeDerivNormal.transpose());
          bfMatrix.noalias() += 0.5 * jacXweightVecSide[j] * (bShapeDerivNormal *  (fShapeMatrix.col(j)).transpose());
          bbMatrix.noalias() += (InputData::penalty/h) * jacXweightVecSide[j] * (bShapeMatrix.col(j) * (bShapeMatrix.col(j)).transpose());
          bfMatrix.noalias() -= (InputData::penalty/h) * jacXweightVecSide[j] * (bShapeMatrix.col(j) * (fShapeMatrix.col(j)).transpose());
        }
        for (unsigned int j = 0; j < InputData::n_dofs; j++)
          for (unsigned int i = 0; i < InputData::n_dofs; i++)
            matrix.insert(elem* InputData::n_dofs+i, neighbor*InputData::n_dofs+j) += bfMatrix(i,j);
      }
    }
    for (unsigned int j = 0; j < InputData::n_dofs; j++)
      for (unsigned int i = 0; i < InputData::n_dofs; i++)
        matrix.insert(elem*InputData::n_dofs+i,elem*InputData::n_dofs+j) += bbMatrix(i,j);
    gRhs.segment<InputData::n_dofs>(elem*InputData::n_dofs).noalias() += rhs;
  }
}
   
   
   
