// SpaFEDTe, a Template based C++ library for creating 
// Discontinuous Finite Element Spaces,
// Copyright (C) 2012 Daniele Di Pietro

/* This library is free software; you can redistribute it and/or */
/* modify it under the terms of the GNU Lesser General Public */ 
/* License as published by the Free Software Foundation either */ 
/* version 3.0 of the License, or (at your option) any later version. */

/* This software is distributed in the hope that it will be useful, */
/* but WITHOUT ANY WARRANTY; without even the implied warranty of */
/* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU */
/* Lesser General Public License for more details. */

/* You should have received a copy of the GNU Lesser General Public */
/* License along with this software; if not, a copy of the full */
/* GNU Lesser General Public License can be found at */
/* http://www.gnu.org/licenses/. */

#include "Debug.hpp"
#include "FELib.hpp"
#include "Print.hpp"

#include "CommonTemplateTypes.hpp"
#include "phoSolution.hpp"
#include "DofMap.hpp"

#include "ReadMeshVtu.hpp"
#include "VTKWriterSpaMesh.hpp"
#include "InputData.hpp"
#include "LocalContributions.hpp"
#include "Timer.hpp"
#include "array"

#include "Eigen/SparseCore"
#include "Eigen/SparseLU"
#include <Eigen/IterativeLinearSolvers>

#define DEBUG_LEVEL 1

// Namespaces
typedef ho::LocalContributions LC;
typedef CommonTemplateTypes::FEManagerType FEManagerType;
typedef CommonTemplateTypes::EigenGlobalVectorType EigenGlobalVectorType;
typedef CommonTemplateTypes::EigenGlobalSolutionType EigenGlobalSolutionType;
typedef CommonTemplateTypes::EigenHOGlobalSolutionType EigenHOGlobalSolutionType;
typedef CommonTemplateTypes::MeshType MeshType;

int main(int argc, char* argv[])
{
  FELib::FELibInit init(argc, argv);
  Timer tmr;
  std::vector<double> l2ErrorNorms;
  std::vector<double> assemblyTimes;
  std::vector<double> solutionTimes;
  std::vector<long unsigned int> matricesNonZeroes;
  std::vector<long unsigned int> nDOFs;

  for (std::string meshFile : InputData::meshSequence)
  {
    // Pretty-print floating point numbers
    std::cout.precision(2);
    std::cout.flags(std::ios_base::scientific);

    // Topology
    MeshType mesh;
    // Load mesh
    ReadMeshVtu vtuReader;
    vtuReader.load_mesh(meshFile.c_str(), mesh);
    //vtuReader.load_mesh("../../meshes_vtu/square_01/square_01_64.vtu", mesh);
    // Compute topology
    bool computeFaces = true;
    mesh.computeElemNeighbors(computeFaces);
    mesh.printInfo();
    std::cout<<"Polynomial degree is "<<InputData::p_order<<", DOFs per face "<<InputData::n_face_dofs<<std::endl;
    
    DofMap dofMap(InputData::n_var, InputData::p_order, InputData::n_face_dofs);
    dofMap.computeMatrixSparsityHHO(mesh); 
    Eigen::VectorXi matrixNonZeroes(dofMap.getGlobalElementNumber()*InputData::n_face_dofs);
    for (unsigned int i = 0; i < dofMap.getGlobalElementNumber(); i++)
      matrixNonZeroes.segment<InputData::n_face_dofs>(i*InputData::n_face_dofs).noalias() = Eigen::VectorXi::Constant(InputData::n_face_dofs, dofMap.getLocalNonZeroBlocks()[i]*InputData::n_face_dofs);
    matricesNonZeroes.push_back(matrixNonZeroes.sum());
    nDOFs.push_back(dofMap.getGlobalElementNumber()*InputData::n_face_dofs);

    //------------------------------------------------------------------------------
    // Exact solution
    ho::Regular sol;
    // ho::pho::Anisotropic sol(1e-1);
    // ho::pho::LePotier sol(1e-2);
    // ho::pho::MixedSecondOrder sol;
    typedef Eigen::SparseMatrix<double> SparseMatrixType;
    typedef Eigen::Matrix<double, Eigen::Dynamic, 1> SolutionVectorType;
    typedef Eigen::Triplet<double> TripletType;
    SparseMatrixType A(dofMap.getGlobalElementNumber()*InputData::n_face_dofs, dofMap.getGlobalElementNumber()*InputData::n_face_dofs);
    std::vector<TripletType> triplets;
    SolutionVectorType X = SolutionVectorType::Zero(dofMap.getGlobalElementNumber()*InputData::n_face_dofs);
    SolutionVectorType B = SolutionVectorType::Zero(dofMap.getGlobalElementNumber()*InputData::n_face_dofs);
   
    tmr.reset();
    LC lc;
    Eigen::VectorXi faces;
    Eigen::Matrix<double, InputData::n_face_dofs, InputData::n_face_dofs, Eigen::RowMajor> AFblock;
    for (SpaMesh::ElementIterator elemIt(mesh); elemIt; ++elemIt) 
    {
      lc.assemble(mesh, elemIt, InputData::p_order, sol.u, sol.nu, sol.f, true);
      faces = mesh.elemMesh.elemSidesConnectivity[elemIt.getType()].col(elemIt.getTypeID());

      for(int i = 0; i < faces.size(); i++) {
        for(int j = 0; j < faces.size(); j++) {
          for(int k = 0; k < InputData::n_face_dofs; k++) {
            for(int p = 0; p < InputData::n_face_dofs; p++) {
              triplets.push_back(TripletType(faces[i]*InputData::n_face_dofs+k, faces[j]*InputData::n_face_dofs+p,
                                 lc.AF(InputData::n_face_dofs*i+k,InputData::n_face_dofs*j+p)));
            }
          }
        } // for i
      } // for j
      for(int i = 0; i < faces.size(); i++) {
        for(int k = 0; k < InputData::n_face_dofs; k++) {
          B(faces[i]*InputData::n_face_dofs+k) += lc.bF(InputData::n_face_dofs*i+k); 
        }
      }
    }
    double assembly_time = tmr.elapsed();
    assemblyTimes.push_back(assembly_time);
    std::cout<<"assembly done"<<std::endl;

    A.setFromTriplets(triplets.begin(), triplets.end());
    tmr.reset();
    Eigen::SparseLU<SparseMatrixType> solver;
    // Compute the ordering permutation vector from the structural pattern of A
    solver.analyzePattern(A);
    // Compute the numerical factorization
    solver.factorize(A);
    //Use the factors to solve the linear system
    X = solver.solve(B);
    //Eigen::ConjugateGradient<SparseMatrixType> solver;
    //solver.setTolerance(1.e-10);
    //solver.compute(A);
    //X = solver.solve(B);
    double solution_time = tmr.elapsed();
    solutionTimes.push_back(solution_time);
    std::cout<<"solve done"<<std::endl;

    const Eigen::Map<const EigenGlobalVectorType> solution(X.data(), mesh.faceCard(), InputData::n_face_dofs); 

    std::cout << FORMAT(50) << "time_assembly"  << assembly_time << std::endl;
    std::cout << FORMAT(50) << "time_solve" << solution_time << std::endl;
    
    //------------------------------------------------------------------------------
    // Compute errors

    double err_GI = 0.;
    double err_Gh = 0.;
    double err_uh = 0.;
    double err_cuh = 0.;
    double norm_u = 0.;
    double norm_G = 0.;
      
    double interp_err = 0.;
    double solL2_err = 0.;
    double solL2_err_ho = 0.;

    const LC::FEManagerType::ShapeMatrixType& shapeMatrix            = lc.FEManager.getElemShapeFunction();
    const LC::FEManagerType::ShapeDerivMatrixType& shapeDerivMatrix  = lc.FEManager.getElemShapeFunctionDeriv();
    const LC::FEManagerType::ElemJacobianVectorType& jacXweightVec   = lc.FEManager.getElemJacobianWeightVec();
    const LC::FEManagerType::ElemQpointVectorType& qpXVec            = lc.FEManager.getElemQpointX();
    
    EigenGlobalSolutionType cellSolution;
    EigenHOGlobalSolutionType cellSolutionHO;
    cellSolution.resize(mesh.elemCard(),Eigen::NoChange);
    cellSolutionHO.resize(mesh.elemCard(),Eigen::NoChange);
    for (SpaMesh::ElementIterator elemIt(mesh); elemIt; ++elemIt) 
    {
      lc.assemble(mesh, elemIt, InputData::p_order, sol.u, sol.nu, sol.f);
      Eigen::VectorXi faces = mesh.elemMesh.elemSidesConnectivity[elemIt.getType()].col(elemIt.getTypeID());

      // Reconstruct local discrete solution
      Eigen::VectorXd uh_TF = lc.reconstruct(solution, faces);

      // Interpolate exact solution
      Eigen::VectorXd u_TF = lc.interpolate(mesh, InputData::p_order, elemIt, sol.u);    
      norm_u += u_TF.head(LC::nb_local_cell_dofs).transpose() * lc.MTT.topLeftCorner(LC::nb_local_cell_dofs, LC::nb_local_cell_dofs) * u_TF.head(LC::nb_local_cell_dofs);

      Eigen::VectorXd diff_T = (uh_TF-u_TF).head(LC::nb_local_cell_dofs);
      err_uh += diff_T.transpose() * lc.MTT.topLeftCorner(LC::nb_local_cell_dofs, LC::nb_local_cell_dofs) * diff_T;

      // High-order potential reconstruction
      Eigen::Matrix<double, LC::NP, 1> cuh_T = lc.potentialReconstruction(uh_TF);

      // Coefficients
      Eigen::Matrix<double, LC::NG, 1> GI_coeffs = lc.GT * u_TF;
      Eigen::Matrix<double, LC::NG, 1> Gh_coeffs = lc.GT * uh_TF;

      unsigned int nElemQP = shapeMatrix.cols();
      for(unsigned int iQN = 0; iQN < nElemQP; iQN++)
      {
        // Iterpolation error
        double u_T_iqn = u_TF.head(LC::nb_cell_dofs).transpose() * shapeMatrix.col(iQN).head(LC::nb_cell_dofs);
        double udiff_T_iqn = u_T_iqn - sol.u(qpXVec.row(iQN));
        interp_err += jacXweightVec[iQN] * std::pow(udiff_T_iqn, 2);
        // solution error
        double uh_T_iqn = uh_TF.head(LC::nb_cell_dofs).transpose() * shapeMatrix.col(iQN).head(LC::nb_cell_dofs);
        double uhdiff_T_iqn = uh_T_iqn - sol.u(qpXVec.row(iQN));
        solL2_err += jacXweightVec[iQN] * std::pow(uhdiff_T_iqn, 2);
        // solution error hhr
        double uh_T_iqn_ho = cuh_T.transpose() * shapeMatrix.col(iQN);
        double uhdiff_T_iqn_ho = uh_T_iqn_ho - sol.u(qpXVec.row(iQN));
        solL2_err_ho += jacXweightVec[iQN] * std::pow(uhdiff_T_iqn_ho, 2);
      } // for iQN
      for(unsigned int iQN = 0; iQN < nElemQP; iQN++)
      {
        Eigen::Matrix<double, InputData::dim, 1> G_iqn  = sol.G(qpXVec.row(iQN));
        Eigen::Matrix<double, InputData::dim, 1> GI_iqn = Eigen::Matrix<double, InputData::dim, 1>::Zero();
        Eigen::Matrix<double, InputData::dim, 1> Gh_iqn = Eigen::Matrix<double, InputData::dim, 1>::Zero();

        for (unsigned int d = 0; d < InputData::dim; d++)
        {
          GI_iqn(d) = GI_coeffs.transpose() * (shapeDerivMatrix(d).block(1,iQN,LC::NG,1)).col(0);
          Gh_iqn(d) = Gh_coeffs.transpose() * (shapeDerivMatrix(d).block(1,iQN,LC::NG,1)).col(0);
        }

        norm_G += jacXweightVec[iQN] * G_iqn.squaredNorm();
        err_GI += jacXweightVec[iQN] * (G_iqn - GI_iqn).dot(G_iqn - GI_iqn);
        err_Gh += jacXweightVec[iQN] * (Gh_iqn - GI_iqn).dot(Gh_iqn - GI_iqn);

        // Displacement reconstruction
        double cuh_T_iqn = cuh_T.transpose() * shapeMatrix.col(iQN);
        double cdiff_T_iqn = cuh_T_iqn - sol.u(qpXVec.row(iQN));
        err_cuh += jacXweightVec[iQN] * std::pow(cdiff_T_iqn, 2);
      } // for iQN

      cellSolution.row(elemIt.value()) = uh_TF.head(LC::nb_cell_dofs);
      cellSolutionHO.row(elemIt.value()) = cuh_T; 
    } // for iT
//  
    err_GI = std::sqrt(err_GI);
    err_uh = std::sqrt(err_uh);
    err_cuh = std::sqrt(err_cuh);
    err_Gh = std::sqrt(err_Gh);
    norm_u = std::sqrt(norm_u);
    norm_G = std::sqrt(norm_G);
      
    interp_err = std::sqrt(interp_err);
    solL2_err = std::sqrt(solL2_err);
    solL2_err_ho = std::sqrt(solL2_err_ho);
    l2ErrorNorms.push_back(solL2_err_ho);

    //std::cout << FORMAT(50) << "err_GI" << err_GI << std::endl;
    //std::cout << FORMAT(50) << "norm_u" << norm_u << std::endl;
    //std::cout << FORMAT(50) << "err_uh" << err_uh / norm_u << std::endl;
    //std::cout << FORMAT(50) << "err_cuh" << err_cuh / norm_u << std::endl;
    //std::cout << FORMAT(50) << "err_Gh" << err_Gh / norm_G << std::endl;
    std::cout << FORMAT(50) << "DOFs" << dofMap.getLocalElementNumber()*InputData::n_face_dofs<< std::endl;
    //std::cout << FORMAT(50) <<"l2 error int"<<interp_err<< std::endl;
    std::cout << FORMAT(50) <<"l2 error "<<solL2_err<< std::endl;
    std::cout << FORMAT(50) <<"l2 error ho"<<solL2_err_ho<< std::endl;

    VTKWriterSpaMesh<InputData::dim, LC::FEManagerType> vtkWriter;
    dofMap.setPOrder(InputData::p_order+1);
    vtkWriter.writeOnTheFly(mesh,lc.FEManager.getFEDataAccessor(),dofMap,cellSolutionHO);
  }
  std::cout<<"DOFs, "<<"Matrix nonZeroes, "<<"Error, "<<"Assembly Time, "<<"Solution Time"<<std::endl;
  for (unsigned int i = 0; i < l2ErrorNorms.size(); i++)
    std::cout<<nDOFs[i]<<","<<matricesNonZeroes[i]<<","<<l2ErrorNorms[i]<<","<<assemblyTimes[i]<<","<<solutionTimes[i]<<std::endl;
  for (unsigned int i = 1; i < l2ErrorNorms.size(); i++)
    std::cout<<"Convergence rate "<<log(l2ErrorNorms[i-1]/l2ErrorNorms[i])/log(2.)<<std::endl;
  return EXIT_SUCCESS;
}
