#include "LocalContributions.hpp"

namespace ho
{
  void LocalContributions::assemble(const MeshType& mesh,
                                    const SpaMesh::ElementIterator& elemIt,
			            const unsigned int p_order,
                                    const ExactSolutionType & exact_solution,
                                    const DiffusivityType & diffusivity,
                                    const LoadType & load,
                                    const bool & weak_bc,
                                    const double & eta)
  {
    int typeDim = elemIt.getType();
    int typeId = elemIt.getTypeID();
    Eigen::VectorXi faces     = mesh.elemMesh.elemSidesConnectivity[typeDim].col(typeId);
    Eigen::VectorXi neighbors = mesh.elemMesh.elemElemsConnectivity[typeDim].col(typeId);
    // Coefficient for the penalty parameter
    std::vector<double> nuF(faces.size(),0.);

    // Count local unknowns
    nb_face_dofs = faces.size() * nb_local_face_dofs;
    nb_tot_dofs = nb_cell_dofs + nb_face_dofs;
    bTF = Eigen::VectorXd::Zero(nb_tot_dofs);

    //------------------------------------------------------------------------------
    // Gradient reconstruction

    Eigen::MatrixXd MG = Eigen::MatrixXd::Zero(NG, NG);
    Eigen::MatrixXd BG = Eigen::MatrixXd::Zero(NG, nb_tot_dofs);
    Eigen::VectorXd bG = Eigen::VectorXd::Zero(NG);
    MTT = Eigen::Matrix<double, NP, NP>::Zero();
    FEManager.faceBasisFunction.doComputeShapeDeriv(false);
    FEManager.basisFunction.doComputeShapeDeriv(true);

    //------------------------------------------------------------------------------
    // Volumetric terms
    {
      const FEManagerType::ShapeMatrixType& shapeMatrix            = FEManager.getElemShapeFunction();
      const FEManagerType::ShapeDerivMatrixType& shapeDerivMatrix  = FEManager.getElemShapeFunctionDeriv();
      const FEManagerType::ElemJacobianVectorType& jacXweightVec   = FEManager.getElemJacobianWeightVec();
      const FEManagerType::ElemQpointVectorType& qpXVec            = FEManager.getElemQpointX();
      FEManager.initVolumeFE(mesh,elemIt,p_order); 

      unsigned int nElemQP = shapeMatrix.cols();
      for(unsigned int iQP = 0; iQP < nElemQP; iQP++) 
      {
        DiffusivityValueType nuT_iqp = diffusivity(qpXVec.row(iQP));

        // LHS
        for (unsigned int di = 0; di < InputData::dim; di++)
          for (unsigned int dj = 0; dj < InputData::dim; dj++)
            MG.noalias() += jacXweightVec[iQP] * nuT_iqp(di,dj) * 
                            (shapeDerivMatrix(di).block(1,iQP,NG,1) * 
        	            (shapeDerivMatrix(dj).block(1,iQP,NG,1)).transpose());

        // RHS (\diff \GRAD vT, \GRAD w)_{PTF}
        for (unsigned int di = 0; di < InputData::dim; di++)
          for (unsigned int dj = 0; dj < InputData::dim; dj++)
            BG.block(0,0,NG,nb_cell_dofs).noalias() += jacXweightVec[iQP] * nuT_iqp(di,dj) * 
                                                       (shapeDerivMatrix(di).block(1,iQP,NG,1) * 
        	                                       (shapeDerivMatrix(dj).block(0,iQP,nb_cell_dofs,1)).transpose());
        // MTT
        MTT.noalias() += jacXweightVec[iQP] * (shapeMatrix.col(iQP) * 
                                                 shapeMatrix.col(iQP).transpose());
        // Forcing term
        bTF.segment(0,nb_cell_dofs).noalias() += jacXweightVec[iQP] * load(qpXVec.row(iQP)) *
                                                  shapeMatrix.block(0,iQP,nb_cell_dofs,1);
      }
    }

    MFF.resize(faces.size());
    std::vector<Eigen::Matrix<double, nb_local_face_dofs, NP> > MFT(faces.size());
    std::vector<Eigen::Matrix<double, nb_local_face_dofs, NG> > NFT(faces.size());

    const FEManagerType::FaceShapeMatrixType& faceShapeMatrix           = FEManager.getFaceShapeFunction();
    const FEManagerType::ShapeMatrixType& elemFaceShapeMatrix           = FEManager.getElemFaceShapeFunction();
    const FEManagerType::ShapeDerivMatrixType& elemFaceShapeDerivMatrix = FEManager.getElemFaceShapeFunctionDeriv();
    const FEManagerType::NormalVectorType& normalVec                    = FEManager.getFaceNormalVec();
    const FEManagerType::FaceJacobianVectorType& jacXweightVecFace      = FEManager.getFaceJacobianWeightVec();
    const FEManagerType::FaceQpointVectorType& qpXVecFace               = FEManager.getFaceQpointX();
    Eigen::Matrix<double, NG, 1> diffusivityShapeDerivNormal_iQN; 

    for(int iF_loc = 0; iF_loc < faces.size(); iF_loc++) 
    {
      FEManager.initFaceElemFaceFE(mesh,elemIt,faces[iF_loc],p_order); 

      //------------------------------------------------------------------------------
      // Initialize matrices
      MFT[iF_loc] = Eigen::Matrix<double, nb_local_face_dofs, NP>::Zero();
      MFF[iF_loc] = Eigen::Matrix<double, nb_local_face_dofs, nb_local_face_dofs>::Zero();
      NFT[iF_loc] = Eigen::Matrix<double, nb_local_face_dofs, NG>::Zero();

      //------------------------------------------------------------------------------
      // Interface terms

      // Compute coefficient for penalty term
      unsigned int nFaceQP = elemFaceShapeMatrix.cols();
      for(int iQN = 0; iQN < nFaceQP; iQN++) 
        nuF[iF_loc] = std::max(nuF[iF_loc], 
                      normalVec[iQN].dot(diffusivity(qpXVecFace.row(iQN)) * normalVec[iQN]));

      for(int iQN = 0; iQN < nFaceQP; iQN++) 
      {
        DiffusivityValueType nuT_iqn = diffusivity(qpXVecFace.row(iQN));

        // Offset for face unknowns
        std::size_t offset_F = nb_cell_dofs + iF_loc * nb_local_face_dofs;

        //Compute \diff\GRAD w\SCAL n_{TF}
        diffusivityShapeDerivNormal_iQN.setZero();
        for (unsigned int di = 0; di < InputData::dim; di++)
          for (unsigned int dj = 0; dj < InputData::dim; dj++)
            diffusivityShapeDerivNormal_iQN.noalias() += normalVec[iQN][di] * nuT_iqn(di,dj) *
                                                           elemFaceShapeDerivMatrix(dj).block(1,iQN,NG,1);

        // RHS (v_F, \diff\GRAD w\SCAL n_{TF})_F
        BG.block(0,offset_F,NG,nb_local_face_dofs).noalias() += jacXweightVecFace[iQN] * 
                                                                  (diffusivityShapeDerivNormal_iQN * 
        							   faceShapeMatrix.col(iQN).transpose());

        // RHS -(v_T, \diff\GRAD w\SCAL n_{TF})_F
        BG.block(0,0,NG,nb_cell_dofs).noalias() -= jacXweightVecFace[iQN] * 
                                                   (diffusivityShapeDerivNormal_iQN * 
                                                    elemFaceShapeMatrix.block(0,iQN,nb_cell_dofs,1).transpose());

        // MFT and NFT
        if(weak_bc)
        {
          const double hF = jacXweightVecFace.sum();
          if (neighbors[iF_loc] == elemIt.value()) 
          {
            bTF.segment(offset_F,nb_local_face_dofs) += (jacXweightVecFace[iQN] * nuF[iF_loc] * 
                                                        (eta / hF) * exact_solution(qpXVecFace.row(iQN))) *
                                                         faceShapeMatrix.col(iQN);
          } // if
        }
        MFT[iF_loc] += jacXweightVecFace[iQN] * (faceShapeMatrix.col(iQN) * 
                                                 elemFaceShapeMatrix.col(iQN).transpose());
        NFT[iF_loc] += jacXweightVecFace[iQN] * (faceShapeMatrix.col(iQN) *
                                                 diffusivityShapeDerivNormal_iQN.transpose());
          
        // MFF
        MFF[iF_loc] += jacXweightVecFace[iQN] * (faceShapeMatrix.col(iQN) * 
                                                 faceShapeMatrix.col(iQN).transpose());
      } // for iQN
    } // for iF_loc

    //------------------------------------------------------------------------------
    // Consistent terms
    
    GT.noalias() = MG.ldlt().solve(BG);

    ATF.noalias() = BG.transpose() * GT;

    //------------------------------------------------------------------------------
    // Stabilization

    STF = Eigen::MatrixXd::Zero(nb_tot_dofs, nb_tot_dofs);
    Eigen::MatrixXd piTK_rTK;
    Eigen::VectorXd rhs_piTK_rTK;
    {
      Eigen::LDLT<Eigen::MatrixXd> piKT;
      piKT.compute(MTT.topLeftCorner(nb_cell_dofs, nb_cell_dofs));

      piTK_rTK = piKT.solve(MTT.block(0, 1, nb_cell_dofs, NG) * GT);
    }

    // Compute face residual
    for(int iF_loc = 0; iF_loc < faces.size(); iF_loc++) 
    {
      const double hF = FEManager.computeFEFaceSize(mesh, faces[iF_loc]);

      Eigen::MatrixXd piF_rTK_minus_uF;
      Eigen::MatrixXd piF_uT_minus_piTK_rTK;
      {
        Eigen::LDLT<Eigen::MatrixXd> piKF;
        piKF.compute(MFF[iF_loc]);

        piF_rTK_minus_uF = piKF.solve(MFT[iF_loc].block(0, 1, nb_local_face_dofs, NG) * GT);
        piF_rTK_minus_uF.block<nb_local_face_dofs, nb_local_face_dofs>(0, nb_cell_dofs + iF_loc * nb_local_face_dofs)
                                              -=  Eigen::Matrix<double, nb_local_face_dofs, nb_local_face_dofs>::Identity();

        Eigen::MatrixXd uT_minus_piTK_rTK = -piTK_rTK;
        uT_minus_piTK_rTK.topLeftCorner(nb_cell_dofs, nb_cell_dofs) += Eigen::MatrixXd::Identity(nb_cell_dofs, nb_cell_dofs);
        piF_uT_minus_piTK_rTK =  piKF.solve(MFT[iF_loc].topLeftCorner(nb_local_face_dofs, nb_cell_dofs) * uT_minus_piTK_rTK);
      }

      auto BRF = piF_uT_minus_piTK_rTK + piF_rTK_minus_uF;

      STF += (nuF[iF_loc] / hF) * BRF.transpose() * MFF[iF_loc] * BRF;
    } // for iF_loc
  
    //------------------------------------------------------------------------------
    // Weakly enforced bc

    if(weak_bc) 
    {
      for(int iF_loc = 0; iF_loc < faces.size(); iF_loc++) 
      {
        const double hF = FEManager.computeFEFaceSize(mesh, faces[iF_loc]);

        std::size_t offset_F = nb_cell_dofs + iF_loc * nb_local_face_dofs;

        if(neighbors[iF_loc] == elemIt.value()) {
          ATF.middleRows<nb_local_face_dofs>(offset_F) -= NFT[iF_loc] * GT;
          STF.block<nb_local_face_dofs, nb_local_face_dofs>(offset_F, offset_F) += (nuF[iF_loc] * eta / hF) * MFF[iF_loc];
        } // if
      } // for
    } // if

    ATF += STF;
  
    //------------------------------------------------------------------------------
    // Static condensation

    Eigen::LDLT<Eigen::MatrixXd> LUATT;
    LUATT.compute(ATF.topLeftCorner(nb_cell_dofs, nb_cell_dofs));

    AF.noalias() = ATF.bottomRightCorner(nb_face_dofs, nb_face_dofs)
                   - ATF.bottomLeftCorner(nb_face_dofs, nb_cell_dofs) * LUATT.solve(ATF.topRightCorner(nb_cell_dofs, nb_face_dofs));

    bF.noalias() = bTF.tail(nb_face_dofs) - ATF.bottomLeftCorner(nb_face_dofs, nb_cell_dofs) * LUATT.solve(bTF.head(nb_cell_dofs));
  } // LocalContributions
  
    //------------------------------------------------------------------------------
  
  Eigen::VectorXd LocalContributions::interpolate(const MeshType& mesh,
      		                                  const unsigned int p_order,
                                                  const SpaMesh::ElementIterator& elemIt,
                                                  const ExactSolutionType & u)
  {
    Eigen::VectorXi faces = mesh.elemMesh.elemSidesConnectivity[elemIt.getType()].col(elemIt.getTypeID());
    Eigen::VectorXd uTF = Eigen::VectorXd::Zero(nb_cell_dofs + faces.size() * nb_local_face_dofs);
    Eigen::VectorXd bT = Eigen::VectorXd::Zero(nb_cell_dofs);
    Eigen::VectorXd bF;
    
    {
      const FEManagerType::ShapeMatrixType& shapeMatrix            = FEManager.getElemShapeFunction();
      const FEManagerType::ElemJacobianVectorType& jacXweightVec   = FEManager.getElemJacobianWeightVec();
      const FEManagerType::ElemQpointVectorType& qpXVec            = FEManager.getElemQpointX();
      FEManager.basisFunction.doComputeShapeDeriv(false);
      FEManager.initVolumeFE(mesh,elemIt,p_order); 
  
      for(int iQN = 0; iQN < shapeMatrix.cols(); iQN++) 
        bT += (jacXweightVec[iQN] * u(qpXVec.row(iQN))) * shapeMatrix.block(0,iQN,nb_cell_dofs,1);
    }

    const FEManagerType::FaceShapeMatrixType& faceShapeMatrix      = FEManager.getFaceShapeFunction();
    const FEManagerType::FaceJacobianVectorType& jacXweightVecFace = FEManager.getFaceJacobianWeightVec();
    const FEManagerType::FaceQpointVectorType& qpXVecFace          = FEManager.getFaceQpointX();
    FEManager.faceBasisFunction.doComputeShapeDeriv(false);
    Eigen::Matrix<double, NG, 1> diffusivityShapeDerivNormal_iQN; 
    for(int iF_loc = 0; iF_loc < faces.size(); iF_loc++) 
    {
      bF = Eigen::VectorXd::Zero(nb_local_face_dofs);
      FEManager.initFaceFE(mesh,elemIt,faces[iF_loc],p_order); 
  
      for(int iQN = 0; iQN < faceShapeMatrix.cols(); iQN++) 
        bF += (jacXweightVecFace[iQN] * u(qpXVecFace.row(iQN))) * faceShapeMatrix.col(iQN);
  
      std::size_t offset_F = nb_cell_dofs + iF_loc * nb_local_face_dofs;
      uTF.segment<nb_local_face_dofs>(offset_F) = MFF[iF_loc].llt().solve(bF);
    } // for iF_loc
  
    Eigen::LDLT<Eigen::MatrixXd> LUMT;
    LUMT.compute(MTT.topLeftCorner(nb_cell_dofs, nb_cell_dofs));
    uTF.head(nb_cell_dofs) = LUMT.solve(bT);
  
    return uTF;
  } // interpolate
  
  //------------------------------------------------------------------------------

  Eigen::VectorXd LocalContributions::reconstruct(const EigenGlobalVectorType & uFh,
                                                  const Eigen::VectorXi & idx_T)
  {
    Eigen::VectorXd uTF = Eigen::VectorXd::Zero(nb_tot_dofs);
    for(int i = 0; i < idx_T.size(); i++) 
    { 
      uTF.segment(nb_cell_dofs+i*nb_local_face_dofs,nb_local_face_dofs) = uFh.row(idx_T[i]); 
    }
    uTF.head(nb_cell_dofs) =
      ATF.topLeftCorner(nb_cell_dofs, nb_cell_dofs).ldlt().solve(
                                                                 bTF.head(nb_cell_dofs)
                                                                 - ATF.topRightCorner(nb_cell_dofs, nb_face_dofs) * uTF.tail(nb_face_dofs)
                                                                 );
    return uTF;
  } // reconstruct
  //------------------------------------------------------------------------------

  Eigen::VectorXd LocalContributions::reconstruct(const Eigen::VectorXd & uTFh)
  {
    Eigen::VectorXd uTF = Eigen::VectorXd::Zero(nb_tot_dofs);
    uTF.head(nb_cell_dofs) =
      ATF.topLeftCorner(nb_cell_dofs, nb_cell_dofs).ldlt().solve(
                                                                 bTF.head(nb_cell_dofs)
                                                                 - ATF.topRightCorner(nb_cell_dofs, nb_face_dofs) * uTFh.tail(nb_face_dofs)
                                                                 );
    return uTF;
  } // reconstruct

  //------------------------------------------------------------------------------
  
  Eigen::Matrix<double, InputData::n_dofs_ppo, 1>
  LocalContributions::potentialReconstruction(const Eigen::VectorXd & uh_TF)
  {
    typedef Eigen::Matrix<double, NP, 1> ResultType;
  
    ResultType cuT = ResultType::Zero();
    cuT.tail(NG) = GT * uh_TF;
    Eigen::Matrix<double, nb_cell_dofs, 1> uT_minus_piKpiuT =
      uh_TF.head<nb_cell_dofs>() - MTT.topLeftCorner(nb_cell_dofs, nb_cell_dofs).ldlt().solve(MTT.topLeftCorner(nb_cell_dofs, NP) * cuT);
  
    cuT.head(nb_cell_dofs) += uT_minus_piKpiuT;
  
    return cuT;
  }
}
